switch to ValueEnum in models

emmet-core/emmet/core/material.py

@@ -24,9 +24,6 @@ class PropertyOrigin(BaseModel):
         default_factory=datetime.utcnow,
     )
 
-    class Config:
-        use_enum_values = True
-
 
 S = TypeVar("S", bound="PropertyDoc")
 

emmet-core/emmet/core/provenance.py

@@ -1,15 +1,15 @@
 """ Core definition of a Provenance Document """
 from datetime import datetime
-from enum import Enum
 from typing import ClassVar, Dict, List
 
 from pybtex.database import BibliographyData, parse_string
 from pydantic import BaseModel, EmailStr, Field, HttpUrl, validator
 
 from emmet.core.material import PropertyDoc
+from emmet.core.utils import ValueEnum
 
 
-class Database(Enum):
+class Database(ValueEnum):
     """
     Database identifiers for provenance IDs
     """

emmet-core/emmet/core/structure.py

@@ -64,9 +64,6 @@ class StructureMetadata(BaseModel):
 
     symmetry: SymmetryData = Field(None, description="Symmetry data for this material")
 
-    class Config:
-        use_enum_values = True
-
     @classmethod
     def from_composition(
         cls: Type[T],

emmet-core/emmet/core/symmetry.py

@@ -1,14 +1,14 @@
-from enum import Enum
 from typing import Any, Dict
 
 from pydantic import BaseModel, Field
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer, spglib
 
 from emmet.core import SETTINGS
+from emmet.core.utils import ValueEnum
 from emmet.stubs import Structure
 
 
-class CrystalSystem(Enum):
+class CrystalSystem(ValueEnum):
     """
     The crystal system of the lattice
     """

emmet-core/emmet/core/vasp/calc_types.py

@@ -1,6 +1,5 @@
 """ Module to define various calculation types as Enums for VASP """
 import datetime
-from enum import Enum
 from itertools import groupby, product
 from pathlib import Path
 from typing import Dict, Iterator, List
@@ -15,6 +14,7 @@
 from typing_extensions import Literal
 
 from emmet.core import SETTINGS
+from emmet.core.utils import ValueEnum
 
 _RUN_TYPE_DATA = loadfn(str(Path(__file__).parent.joinpath("run_types.yaml").resolve()))
 _TASK_TYPES = [
@@ -44,15 +44,15 @@
     + ["LDA", "LDA+U"]
 )
 
-RunType = Enum(  # type: ignore
+RunType = ValueEnum(  # type: ignore
     "RunType", dict({"_".join(rt.split()).replace("+", "_"): rt for rt in _RUN_TYPES})
 )
 RunType.__doc__ = "VASP calculation run types"
 
-TaskType = Enum("TaskType", {"_".join(tt.split()): tt for tt in _TASK_TYPES})  # type: ignore
+TaskType = ValueEnum("TaskType", {"_".join(tt.split()): tt for tt in _TASK_TYPES})  # type: ignore
 TaskType.__doc__ = "VASP calculation task types"
 
-CalcType = Enum(  # type: ignore
+CalcType = ValueEnum(  # type: ignore
     "CalcType",
     {
         f"{'_'.join(rt.split()).replace('+','_')}_{'_'.join(tt.split())}": f"{rt} {tt}"

emmet-core/emmet/core/vasp/task.py

@@ -1,6 +1,5 @@
 """ Core definition of a VASP Task Document """
 from datetime import datetime
-from enum import Enum
 from functools import lru_cache, partial
 from typing import ClassVar, Dict, List, Optional, Union
 
@@ -10,6 +9,7 @@
 
 from emmet.core import SETTINGS
 from emmet.core.structure import StructureMetadata
+from emmet.core.utils import ValueEnum
 from emmet.core.vasp.calc_types import (
     CalcType,
     RunType,
@@ -21,7 +21,7 @@
 from emmet.stubs import ComputedEntry, Matrix3D, Structure, Vector3D
 
 
-class Status(Enum):
+class Status(ValueEnum):
     """
     VASP Calculation State
     """

emmet-core/emmet/core/vasp/validation.py

@@ -1,5 +1,4 @@
 from datetime import datetime
-from enum import Enum
 from typing import Dict, List, Union
 
 import numpy as np

emmet-core/emmet/core/xas.py

@@ -1,11 +1,11 @@
 from datetime import datetime
-from enum import Enum
 from typing import List, Optional
 
 from pydantic import BaseModel, Field, root_validator
 from pymatgen.core.periodic_table import Element
 
 from emmet.core.spectrum import SpectrumDoc
+from emmet.core.utils import ValueEnum
 from emmet.stubs import XAS, Structure
 from emmet.stubs.xas import Edge, Type
 

emmet-core/emmet/stubs/xas.py

@@ -1,13 +1,13 @@
-from enum import Enum
 from typing import Dict, List, Optional, Tuple, Union
 
 from pydantic import BaseModel, Field
 from pymatgen.core.periodic_table import Element
 
+from emmet.core.utils import ValueEnum
 from emmet.stubs import Structure
 
 
-class Edge(Enum):
+class Edge(ValueEnum):
     """
     The interaction edge for XAS
     There are 2n-1 sub-components to each edge where
@@ -23,7 +23,7 @@ class Edge(Enum):
     L2_3 = "L2,3"
 
 
-class Type(Enum):
+class Type(ValueEnum):
     """
     The type of XAS Spectrum
     XANES - Just the near-edge region
@@ -54,6 +54,3 @@ class XASSpectrum(BaseModel):
     edge: Edge = Field(
         None, title="Absorption Edge", description="The interaction edge for XAS"
     )
-
-    class Config:
-        use_enum_values = True