diff --git a/emmet-core/emmet/core/common.py b/emmet-core/emmet/core/common.py index 8925e13ebb..63e0f182f0 100644 --- a/emmet-core/emmet/core/common.py +++ b/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. """ diff --git a/emmet-core/emmet/core/electrode.py b/emmet-core/emmet/core/electrode.py index f1a220d974..26f3c5bd1c 100644 --- a/emmet-core/emmet/core/electrode.py +++ b/emmet-core/emmet/core/electrode.py @@ -1,16 +1,15 @@ 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 @@ -18,7 +17,7 @@ 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) diff --git a/emmet-core/emmet/core/feff/task.py b/emmet-core/emmet/core/feff/task.py index ab9881ab02..30c8733b8e 100644 --- a/emmet-core/emmet/core/feff/task.py +++ b/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" ) diff --git a/emmet-core/emmet/core/provenance.py b/emmet-core/emmet/core/provenance.py index 85b564a7d1..77628bd25d 100644 --- a/emmet-core/emmet/core/provenance.py +++ b/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] diff --git a/emmet-core/emmet/core/qchem/calc_types/enums.py b/emmet-core/emmet/core/qchem/calc_types/enums.py index d7205610a5..9bcb515446 100644 --- a/emmet-core/emmet/core/qchem/calc_types/enums.py +++ b/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" diff --git a/emmet-core/emmet/core/qchem/calc_types/generate.py b/emmet-core/emmet/core/qchem/calc_types/generate.py index fd20cf1d2c..23e0a32595 100644 --- a/emmet-core/emmet/core/qchem/calc_types/generate.py +++ b/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 " @@ -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 """ ) diff --git a/emmet-core/emmet/core/qchem/calculation.py b/emmet-core/emmet/core/qchem/calculation.py index f4af5d1484..57578e4f49 100644 --- a/emmet-core/emmet/core/qchem/calculation.py +++ b/emmet-core/emmet/core/qchem/calculation.py @@ -3,26 +3,26 @@ # 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 ( @@ -30,7 +30,6 @@ # task_type, # calc_type, # ) - from emmet.core.qchem.task import QChemStatus functional_synonyms = { @@ -50,7 +49,7 @@ __author__ = "Rishabh D. Guha " 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 diff --git a/emmet-core/emmet/core/qchem/task.py b/emmet-core/emmet/core/qchem/task.py index d009208f08..2a87f1f437 100644 --- a/emmet-core/emmet/core/qchem/task.py +++ b/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 " -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 diff --git a/emmet-core/emmet/core/symmetry.py b/emmet-core/emmet/core/symmetry.py index 4c27155efb..a0d4d3e06a 100644 --- a/emmet-core/emmet/core/symmetry.py +++ b/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 diff --git a/emmet-core/emmet/core/thermo.py b/emmet-core/emmet/core/thermo.py index e1ff775da3..a888b9a7d3 100644 --- a/emmet-core/emmet/core/thermo.py +++ b/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 diff --git a/emmet-core/emmet/core/utils.py b/emmet-core/emmet/core/utils.py index 53762d4c5d..7d48ff4078 100644 --- a/emmet-core/emmet/core/utils.py +++ b/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()] diff --git a/emmet-core/emmet/core/vasp/calc_types/enums.py b/emmet-core/emmet/core/vasp/calc_types/enums.py index 5b28246454..dc469a4003 100644 --- a/emmet-core/emmet/core/vasp/calc_types/enums.py +++ b/emmet-core/emmet/core/vasp/calc_types/enums.py @@ -3,10 +3,10 @@ Do not edit this by hand. Edit generate.py or run_types.yaml instead. """ -from emmet.core.utils import ValueEnum +from enum import StrEnum -class RunType(ValueEnum): +class RunType(Str): """VASP calculation run types.""" AM05 = "AM05" @@ -81,9 +81,10 @@ def _missing_(cls, value): for member in cls: if member.value.upper() == value.upper(): return member + return None -class TaskType(ValueEnum): +class TaskType(StrEnum): """VASP calculation task types.""" NSCF_Line = "NSCF Line" @@ -101,7 +102,7 @@ class TaskType(ValueEnum): Unrecognized = "Unrecognized" -class CalcType(ValueEnum): +class CalcType(StrEnum): """VASP calculation types.""" AM05_NSCF_Line = "AM05 NSCF Line" diff --git a/emmet-core/emmet/core/vasp/calc_types/generate.py b/emmet-core/emmet/core/vasp/calc_types/generate.py index 30b8335a3a..09a57e4dc8 100644 --- a/emmet-core/emmet/core/vasp/calc_types/generate.py +++ b/emmet-core/emmet/core/vasp/calc_types/generate.py @@ -66,7 +66,7 @@ Do not edit this by hand. Edit generate.py or run_types.yaml instead. \"\"\" -from emmet.core.utils import ValueEnum +from enum import StrEnum """ ) diff --git a/emmet-core/emmet/core/vasp/calculation.py b/emmet-core/emmet/core/vasp/calculation.py index 94295fde2c..530611aa7d 100644 --- a/emmet-core/emmet/core/vasp/calculation.py +++ b/emmet-core/emmet/core/vasp/calculation.py @@ -4,6 +4,7 @@ import logging from datetime import datetime +from enum import StrEnum from pathlib import Path from typing import Any, Dict, List, Optional, Tuple, Union @@ -31,7 +32,6 @@ ) from emmet.core.math import ListMatrix3D, Matrix3D, Vector3D -from emmet.core.utils import ValueEnum from emmet.core.vasp.calc_types import ( CalcType, RunType, @@ -45,7 +45,7 @@ logger = logging.getLogger(__name__) -class VaspObject(ValueEnum): +class VaspObject(StrEnum): """Types of VASP data objects.""" BANDSTRUCTURE = "bandstructure" @@ -62,7 +62,7 @@ class VaspObject(ValueEnum): PROCAR = "procar" -class StoreTrajectoryOption(ValueEnum): +class StoreTrajectoryOption(StrEnum): FULL = "full" PARTIAL = "partial" NO = "no" diff --git a/emmet-core/emmet/core/vasp/task_valid.py b/emmet-core/emmet/core/vasp/task_valid.py index 9a95688060..290840bbe9 100644 --- a/emmet-core/emmet/core/vasp/task_valid.py +++ b/emmet-core/emmet/core/vasp/task_valid.py @@ -1,7 +1,8 @@ # mypy: ignore-errors """ Core definition of a VASP Task Document """ -from typing import Any, Dict, List, Union, Optional +from enum import StrEnum +from typing import Any, Dict, List, Optional, Union from pydantic import BaseModel, Field from pymatgen.analysis.structure_analyzer import oxide_type @@ -9,13 +10,12 @@ from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry from emmet.core.math import Matrix3D, Vector3D -from emmet.core.task import BaseTaskDocument from emmet.core.structure import StructureMetadata -from emmet.core.utils import ValueEnum +from emmet.core.task import BaseTaskDocument from emmet.core.vasp.calc_types import RunType, calc_type, run_type, task_type -class TaskState(ValueEnum): +class TaskState(StrEnum): """ VASP Calculation State """ diff --git a/emmet-core/emmet/core/xas.py b/emmet-core/emmet/core/xas.py index f9f0bc54f0..2b52b0c9c8 100644 --- a/emmet-core/emmet/core/xas.py +++ b/emmet-core/emmet/core/xas.py @@ -1,4 +1,5 @@ import warnings +from enum import StrEnum from itertools import groupby from typing import List @@ -11,10 +12,9 @@ from emmet.core.feff.task import TaskDocument from emmet.core.mpid import MPID from emmet.core.spectrum import SpectrumDoc -from emmet.core.utils import ValueEnum -class Edge(ValueEnum): +class Edge(StrEnum): """ The interaction edge for XAS There are 2n-1 sub-components to each edge where @@ -30,7 +30,7 @@ class Edge(ValueEnum): L2_3 = "L2,3" -class Type(ValueEnum): +class Type(StrEnum): """ The type of XAS Spectrum XANES - Just the near-edge region diff --git a/emmet-core/emmet/core/xrd.py b/emmet-core/emmet/core/xrd.py index 67fc4f2b1d..c52a58879e 100644 --- a/emmet-core/emmet/core/xrd.py +++ b/emmet-core/emmet/core/xrd.py @@ -1,7 +1,8 @@ +from enum import StrEnum from typing import Dict, Optional import numpy as np -from pydantic import model_validator, Field +from pydantic import Field, model_validator from pymatgen.analysis.diffraction.xrd import ( WAVELENGTHS, DiffractionPattern, @@ -12,10 +13,9 @@ from emmet.core.mpid import MPID from emmet.core.spectrum import SpectrumDoc -from emmet.core.utils import ValueEnum -class Edge(ValueEnum): +class Edge(StrEnum): K_Alpha = "Ka" K_Alpha1 = "Ka1" K_Alpha2 = "Ka2"