Replace emmet.core.utils.ValueEnum with std lib enum.StrEnum

emmet-core/emmet/core/common.py

@@ -1,7 +1,7 @@
-from emmet.core.utils import ValueEnum
+from enum import StrEnum
 
 
-class Status(ValueEnum):
+class Status(StrEnum):
     """
     State of a calculation/analysis.
     """

emmet-core/emmet/core/electrode.py

@@ -1,24 +1,23 @@
 import re
-from datetime import datetime
-from typing import List, Union, Dict, Optional
 from collections import defaultdict
-
-from emmet.core.utils import ValueEnum
+from datetime import datetime
+from enum import StrEnum
+from typing import Dict, List, Optional, Union
 
 from monty.json import MontyDecoder
-from pydantic import field_validator, BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
+from pymatgen.analysis.phase_diagram import PhaseDiagram
 from pymatgen.apps.battery.battery_abc import AbstractElectrode
 from pymatgen.apps.battery.conversion_battery import ConversionElectrode
 from pymatgen.apps.battery.insertion_battery import InsertionElectrode
-from pymatgen.analysis.phase_diagram import PhaseDiagram
 from pymatgen.core import Composition, Structure
 from pymatgen.core.periodic_table import Element
 from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
 
 from emmet.core.mpid import MPID
 
 
-class BatteryType(str, ValueEnum):
+class BatteryType(StrEnum):
     """
     Enum for battery type
     """
@@ -323,7 +322,7 @@ def from_entries(
 
         # Check if more than one working ion per transition metal and warn
         warnings = []
-        if any([element.is_transition_metal for element in dchg_comp]):
+        if any(element.is_transition_metal for element in dchg_comp):
             transition_metal_fraction = sum(
                 [
                     dchg_comp.get_atomic_fraction(elem)

emmet-core/emmet/core/feff/task.py

@@ -1,4 +1,5 @@
 """ Core definition of a VASP Task Document """
+from enum import StrEnum
 from typing import Any, Dict, List
 
 from pydantic import Field
@@ -8,10 +9,9 @@
 
 from emmet.core.structure import StructureMetadata
 from emmet.core.vasp.task_valid import TaskDocument as BaseTaskDocument
-from emmet.core.utils import ValueEnum
 
 
-class CalcType(ValueEnum):
+class CalcType(StrEnum):
     """
     The type of FEFF Calculation
     XANES - Just the near-edge region
@@ -56,7 +56,7 @@ def absorbing_element(self) -> Element:
     @property
     def xas_spectrum(self) -> XAS:
         if not hasattr(self, "_xas_spectrum"):
-            if not all([len(p) == 6 for p in self.spectrum]):
+            if not all(len(p) == 6 for p in self.spectrum):
                 raise ValueError(
                     "Spectrum data doesn't appear to be from xmu.dat which holds XAS data"
                 )

emmet-core/emmet/core/provenance.py

@@ -1,19 +1,19 @@
 """ Core definition of a Provenance Document """
 import warnings
 from datetime import datetime
+from enum import StrEnum
 from typing import Dict, List, Optional
 
 from pybtex.database import BibliographyData, parse_string
 from pybtex.errors import set_strict_mode
-from pydantic import field_validator, model_validator, BaseModel, Field
+from pydantic import BaseModel, Field, field_validator, model_validator
+from pymatgen.core.structure import Structure
 
 from emmet.core.material_property import PropertyDoc
 from emmet.core.mpid import MPID
-from emmet.core.utils import ValueEnum
-from pymatgen.core.structure import Structure
 
 
-class Database(ValueEnum):
+class Database(StrEnum):
     """
     Database identifiers for provenance IDs
     """
@@ -179,7 +179,7 @@ def from_SNLs(
 
         # TODO: Maybe we should combine this robocrystallographer?
         # TODO: Refine these tags / remarks
-        remarks = list(set([remark for snl in snls for remark in snl.about.remarks]))
+        remarks = list({remark for snl in snls for remark in snl.about.remarks})
         tags = [r for r in remarks if len(r) < 140]
 
         authors = [entry for snl in snls for entry in snl.about.authors]

emmet-core/emmet/core/qchem/calc_types/enums.py

@@ -2,10 +2,10 @@
 Autogenerated Enums for Q-Chem LevelOfTheory, TaskType, and CalcType
 Do not edit this by hand. Edit generate.py or types.py instead
 """
-from emmet.core.utils import ValueEnum
+from enum import StrEnum
 
 
-class LevelOfTheory(ValueEnum):
+class LevelOfTheory(StrEnum):
     """Levels of theory for calculations in Q-Chem"""
 
     PBE_6_31g_d_VACUUM = "PBE/6-31g*/VACUUM"
@@ -250,7 +250,7 @@ class LevelOfTheory(ValueEnum):
     wB97M_V_def2_QZVPPD_SMD = "wB97M-V/def2-QZVPPD/SMD"
 
 
-class TaskType(ValueEnum):
+class TaskType(StrEnum):
     """Calculation task types for Q-Chem"""
 
     Single_Point = "Single Point"
@@ -267,7 +267,7 @@ class TaskType(ValueEnum):
     Unknown = "Unknown"
 
 
-class CalcType(ValueEnum):
+class CalcType(StrEnum):
     """Calculation types (LOT + task type) for Q-Chem"""
 
     PBE_6_31g_d_VACUUM_Single_Point = "PBE/6-31g*/VACUUM Single Point"

emmet-core/emmet/core/qchem/calc_types/generate.py

@@ -2,14 +2,13 @@
 from itertools import product
 from pathlib import Path
 
-from emmet.core.utils import get_enum_source
 from emmet.core.qchem.calc_types.calc_types import (
-    TASK_TYPES,
-    FUNCTIONALS,
     BASIS_SETS,
+    FUNCTIONALS,
     SOLVENT_MODELS,
+    TASK_TYPES,
 )
-
+from emmet.core.utils import get_enum_source
 
 __author__ = "Evan Spotte-Smith <[email protected]>"
 
@@ -48,7 +47,7 @@
     "CalcType",
     "Calculation types (LOT + task type) for Q-Chem",
     {
-        f"{'_'.join(lot.split()).replace('+','_').replace('-','_').replace('(', '_').replace(')', '_').replace('/', '_').replace('*', '_d')}_{'_'.join(tt.split()).replace('-', '_')}": f"{lot} {tt}"  # noqa: E501
+        f"{'_'.join(lot.split()).replace('+','_').replace('-','_').replace('(', '_').replace(')', '_').replace('/', '_').replace('*', '_d')}_{'_'.join(tt.split()).replace('-', '_')}": f"{lot} {tt}"
         for lot, tt in product(_LOTS, TASK_TYPES)
     },
 )
@@ -59,7 +58,7 @@
 Autogenerated Enums for Q-Chem LevelOfTheory, TaskType, and CalcType
 Do not edit this by hand. Edit generate.py or types.py instead
 \"\"\"
-from emmet.core.utils import ValueEnum
+from enum import StrEnum
 
 """
     )

emmet-core/emmet/core/qchem/calculation.py

@@ -3,34 +3,33 @@
 # mypy: ignore-errors
 
 import logging
+import re
+from collections import OrderedDict
+from datetime import datetime
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
-from pydantic import field_validator, BaseModel, Field, ConfigDict
-from datetime import datetime
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pymatgen.core.structure import Molecule
 from pymatgen.io.qchem.inputs import QCInput
 from pymatgen.io.qchem.outputs import QCOutput
-from pymatgen.core.structure import Molecule
-from collections import OrderedDict
-import re
 
 from emmet.core.qchem.calc_types import (
-    LevelOfTheory,
     CalcType,
+    LevelOfTheory,
     TaskType,
 )
 from emmet.core.qchem.calc_types.calc_types import (
-    FUNCTIONALS,
     BASIS_SETS,
+    FUNCTIONALS,
 )
 
 # from emmet.core.qchem.calc_types.em_utils import (
 #     level_of_theory,
 #     task_type,
 #     calc_type,
 # )
-
 from emmet.core.qchem.task import QChemStatus
 
 functional_synonyms = {
@@ -50,7 +49,7 @@
 __author__ = "Rishabh D. Guha <[email protected]>"
 logger = logging.getLogger(__name__)
 
-# class QChemObject(ValueEnum):
+# class QChemObject(StrEnum):
 # Not sure but can we have something like GRAD and HESS
 # as QChem data objects
 

emmet-core/emmet/core/qchem/task.py

@@ -1,30 +1,29 @@
 # mypy: ignore-errors
 
 """ Core definition of a Q-Chem Task Document """
-from typing import Any, Dict, List, Optional, Callable
+from enum import StrEnum
+from typing import Any, Callable, Dict, List, Optional
 
 from pydantic import BaseModel, Field
 from pymatgen.core.structure import Molecule
 
-from emmet.core.structure import MoleculeMetadata
-from emmet.core.task import BaseTaskDocument
-from emmet.core.utils import ValueEnum
 from emmet.core.qchem.calc_types import (
-    LevelOfTheory,
     CalcType,
+    LevelOfTheory,
     TaskType,
     calc_type,
     level_of_theory,
-    task_type,
-    solvent,
     lot_solvent_string,
+    solvent,
+    task_type,
 )
-
+from emmet.core.structure import MoleculeMetadata
+from emmet.core.task import BaseTaskDocument
 
 __author__ = "Evan Spotte-Smith <[email protected]>"
 
 
-class QChemStatus(ValueEnum):
+class QChemStatus(StrEnum):
     """
     Q-Chem Calculation State
     """
@@ -177,10 +176,7 @@ def entry(self) -> Dict[str, Any]:
         else:
             mol = self.output.initial_molecule
 
-        if self.charge is None:
-            charge = int(mol.charge)
-        else:
-            charge = int(self.charge)
+        charge = int(mol.charge) if self.charge is None else int(self.charge)
 
         if self.spin_multiplicity is None:
             spin = mol.spin_multiplicity

emmet-core/emmet/core/symmetry.py

@@ -1,3 +1,4 @@
+from enum import StrEnum
 from typing import Any, Dict, Optional
 
 from pydantic import BaseModel, Field
@@ -6,12 +7,11 @@
 from pymatgen.symmetry.analyzer import PointGroupAnalyzer, SpacegroupAnalyzer, spglib
 
 from emmet.core.settings import EmmetSettings
-from emmet.core.utils import ValueEnum
 
 SETTINGS = EmmetSettings()
 
 
-class CrystalSystem(ValueEnum):
+class CrystalSystem(StrEnum):
     """
     The crystal system of the lattice
     """
@@ -100,10 +100,7 @@ def from_molecule(cls, molecule: Molecule) -> "PointGroupData":
             if symmetry["point_group"] in point_groups:
                 r = rot_num
                 break
-        if symmetry["point_group"] in ["C*v", "D*h"]:
-            linear = True
-        else:
-            linear = False
+        linear = symmetry["point_group"] in ["C*v", "D*h"]
 
         symmetry["rotation_number"] = float(r)
         symmetry["linear"] = linear

emmet-core/emmet/core/thermo.py

@@ -1,16 +1,16 @@
 """ Core definition of a Thermo Document """
 from collections import defaultdict
-from typing import Dict, List, Optional, Union
 from datetime import datetime
-from emmet.core.base import EmmetMeta
-from emmet.core.utils import ValueEnum
+from enum import StrEnum
+from typing import Dict, List, Optional, Union
 
 from pydantic import BaseModel, Field
 from pymatgen.analysis.phase_diagram import PhaseDiagram
 from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
 
-from emmet.core.material_property import PropertyDoc
+from emmet.core.base import EmmetMeta
 from emmet.core.material import PropertyOrigin
+from emmet.core.material_property import PropertyDoc
 from emmet.core.mpid import MPID
 from emmet.core.vasp.calc_types.enums import RunType
 
@@ -34,7 +34,7 @@ class DecompositionProduct(BaseModel):
     )
 
 
-class ThermoType(ValueEnum):
+class ThermoType(StrEnum):
     GGA_GGA_U = "GGA_GGA+U"
     GGA_GGA_U_R2SCAN = "GGA_GGA+U_R2SCAN"
     R2SCAN = "R2SCAN"
@@ -125,7 +125,7 @@ def from_entries(
         thermo_type: Union[ThermoType, RunType],
         phase_diagram: Optional[PhaseDiagram] = None,
         use_max_chemsys: bool = False,
-        **kwargs
+        **kwargs,
     ):
         """Produce a list of ThermoDocs from a list of Entry objects
 

emmet-core/emmet/core/utils.py

@@ -5,11 +5,8 @@
 from typing import Any, Dict, Iterator, List, Optional, Union
 
 import numpy as np
-
 from monty.json import MSONable
-
 from pydantic import BaseModel
-
 from pymatgen.analysis.graphs import MoleculeGraph
 from pymatgen.analysis.local_env import OpenBabelNN, metal_edge_extender
 from pymatgen.analysis.molecule_matcher import MoleculeMatcher
@@ -105,7 +102,7 @@ def _mol_form(mol_solv):
 
     # First, group by formula
     # Hopefully this step is unnecessary - builders should already be doing this
-    for mol_key, pregroup in groupby(sorted(molecules, key=_mol_form), key=_mol_form):
+    for _mol_key, pregroup in groupby(sorted(molecules, key=_mol_form), key=_mol_form):
         groups: List[Dict[str, Any]] = list()
         for mol in pregroup:
             mol_copy = copy.deepcopy(mol)
@@ -334,7 +331,7 @@ def __new__(cls, value, doc=None):
 
 def get_enum_source(enum_name, doc, items):
     header = f"""
-class {enum_name}(ValueEnum):
+class {enum_name}(StrEnum):
     \"\"\" {doc} \"\"\"\n
 """
     items = [f'    {const} = "{val}"' for const, val in items.items()]