Add ValueEnum

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

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"