From e7ba9e2a1729387dfc576ded93f4b402f1c46c12 Mon Sep 17 00:00:00 2001
From: Shyam D <shyamd@lbl.gov>
Date: Fri, 23 Oct 2020 14:20:37 -0700
Subject: [PATCH] Add ValueEnum

---
 emmet-core/emmet/core/utils.py | 11 ++++++++++-
 tests/emmet-core/test_utils.py | 22 ++++++++++++++++++++--
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/emmet-core/emmet/core/utils.py b/emmet-core/emmet/core/utils.py
index 753aea4ea6..68e887ef96 100644
--- a/emmet-core/emmet/core/utils.py
+++ b/emmet-core/emmet/core/utils.py
@@ -125,7 +125,16 @@ def jsanitize(obj, strict=False, allow_bson=False):
     return jsanitize(obj.as_dict(), strict=strict, allow_bson=allow_bson)
 
 
-class DocEnum(Enum):
+class ValueEnum(Enum):
+    """
+    Enum that serializes to string as the value
+    """
+
+    def __str__(self):
+        return str(self.value)
+
+
+class DocEnum(ValueEnum):
     """
     Enum with docstrings support
     from: https://stackoverflow.com/a/50473952
diff --git a/tests/emmet-core/test_utils.py b/tests/emmet-core/test_utils.py
index 100d495783..868d5c4292 100644
--- a/tests/emmet-core/test_utils.py
+++ b/tests/emmet-core/test_utils.py
@@ -2,14 +2,14 @@
 import json
 import os
 import unittest
-from enum import Enum
+
 
 import numpy as np
 import pytest
 from bson.objectid import ObjectId
 from monty.json import MSONable
 
-from emmet.core.utils import get_sg, group_structures, jsanitize
+from emmet.core.utils import get_sg, group_structures, jsanitize, ValueEnum, DocEnum
 
 
 def test_jsanitize():
@@ -69,3 +69,21 @@ def __eq__(self, other):
             and self._d == other._d
             and self.kwargs == other.kwargs
         )
+
+
+def test_value_enum():
+    class TempEnum(ValueEnum):
+        A = "A"
+        B = "B"
+
+    assert str(TempEnum.A) == "A"
+    assert str(TempEnum.B) == "B"
+
+
+def test_doc_enum():
+    class TestEnum(DocEnum):
+        A = "A", "Describes A"
+        B = "B", "Might describe B"
+
+    assert str(TestEnum.A) == "A"
+    assert TestEnum.B.__doc__ == "Might describe B"