[WIP] Strict typing and fixed schemas for emmet document models

emmet-api/tests/materials/tasks/test_utils.py

@@ -1,5 +1,6 @@
 import os
 from json import load
+from monty.io import zopen
 
 from emmet.api.core.settings import MAPISettings
 from emmet.api.routes.materials.tasks.utils import (
@@ -24,7 +25,7 @@ def test_calcs_reversed_to_trajectory():
 
 
 def test_task_to_entry():
-    with open(os.path.join(MAPISettings().TEST_FILES, "test_task.json")) as file:
+    with zopen(os.path.join(MAPISettings().TEST_FILES, "test_task.json.gz")) as file:
         doc = load(file)
 
     entry_dict = task_to_entry(doc)

emmet-builders/emmet/builders/materials/corrected_entries.py

@@ -143,7 +143,6 @@ def process_item(self, item):
         self.logger.debug(f"Processing {len(entries)} entries for {chemsys}")
 
         all_entry_types = {str(e.data["run_type"]) for e in entries}
-
         corrected_entries = {}
 
         for compatibility in self.compatibility:

emmet-core/emmet/core/math.py

@@ -1,25 +1,25 @@
-from typing import Tuple, List
+"""Define types used for array-like data."""
 
-Vector3D = Tuple[float, float, float]
-Vector3D.__doc__ = "Real space vector"  # type: ignore
+Vector3D = tuple[float, float, float]
+"""Real space vector."""
 
-Matrix3D = Tuple[Vector3D, Vector3D, Vector3D]
-Matrix3D.__doc__ = "Real space Matrix"  # type: ignore
+Matrix3D = tuple[Vector3D, Vector3D, Vector3D]
+"""Real space Matrix."""
 
-Vector6D = Tuple[float, float, float, float, float, float]
-Vector6D.__doc__ = "6D Voigt matrix component"  # type: ignore
+Vector6D = tuple[float, float, float, float, float, float]
+"""6D Voigt matrix component."""
 
-MatrixVoigt = Tuple[Vector6D, Vector6D, Vector6D, Vector6D, Vector6D, Vector6D]
-Vector6D.__doc__ = "Voigt representation of a 3x3x3x3 tensor"  # type: ignore
+MatrixVoigt = tuple[Vector6D, Vector6D, Vector6D, Vector6D, Vector6D, Vector6D]
+""""Voigt representation of a 3x3x3x3 tensor."""
 
-Tensor3R = List[List[List[float]]]
-Tensor3R.__doc__ = "Generic tensor of rank 3"  # type: ignore
+Tensor3R = list[list[list[float]]]
+"""Generic tensor of rank 3."""
 
-Tensor4R = List[List[List[List[float]]]]
-Tensor4R.__doc__ = "Generic tensor of rank 4"  # type: ignore
+Tensor4R = list[list[list[list[float]]]]
+"""Generic tensor of rank 4."""
 
-ListVector3D = List[float]
-ListVector3D.__doc__ = "Real space vector as list"  # type: ignore
+ListVector3D = list[float]
+"""Real space vector as list."""
 
-ListMatrix3D = List[ListVector3D]
-ListMatrix3D.__doc__ = "Real space Matrix as list"  # type: ignore
+ListMatrix3D = list[ListVector3D]
+"""Real space Matrix as list."""

emmet-core/emmet/core/structure.py

@@ -1,7 +1,7 @@
 """Core definition of Structure and Molecule metadata."""
 from __future__ import annotations
 
-from typing import List, Optional, Type, TypeVar
+from typing import Optional, Type, TypeVar
 
 from pydantic import Field
 from pymatgen.core.composition import Composition
@@ -11,8 +11,9 @@
 from emmet.core.utils import get_graph_hash
 from emmet.core.base import EmmetBaseModel
 from emmet.core.symmetry import PointGroupData, SymmetryData
+from emmet.core.structure_replicas import CompositionReplica, StructureReplica
 
-T = TypeVar("T", bound="StructureMetadata")
+T = TypeVar("T", bound="DbStructureMetadata")
 S = TypeVar("S", bound="MoleculeMetadata")
 
 try:
@@ -21,21 +22,21 @@
     openbabel = None
 
 
-class StructureMetadata(EmmetBaseModel):
+class DbStructureMetadata(EmmetBaseModel):
     """Mix-in class for structure metadata."""
 
     # Structure metadata
     nsites: Optional[int] = Field(
         None, description="Total number of sites in the structure."
     )
-    elements: Optional[List[Element]] = Field(
+    elements: Optional[list[Element]] = Field(
         None, description="List of elements in the material."
     )
     nelements: Optional[int] = Field(None, description="Number of elements.")
-    composition: Optional[Composition] = Field(
+    composition: Optional[CompositionReplica] = Field(
         None, description="Full composition for the material."
     )
-    composition_reduced: Optional[Composition] = Field(
+    composition_reduced: Optional[CompositionReplica] = Field(
         None,
         title="Reduced Composition",
         description="Simplified representation of the composition.",
@@ -75,11 +76,13 @@ class StructureMetadata(EmmetBaseModel):
         None, description="Symmetry data for this material."
     )
 
+    _use_pymatgen_rep: bool = False
+
     @classmethod
     def from_composition(
         cls: Type[T],
-        composition: Composition,
-        fields: Optional[List[str]] = None,
+        composition: Composition | CompositionReplica,
+        fields: Optional[list[str]] = None,
         **kwargs,
     ) -> T:
         fields = (
@@ -95,6 +98,11 @@ def from_composition(
             if fields is None
             else fields
         )
+        if isinstance(composition, Composition) and not cls._use_pymatgen_rep:
+            composition = CompositionReplica.from_pymatgen(composition)
+        elif isinstance(composition, CompositionReplica) and cls._use_pymatgen_rep:
+            composition = composition.to_pymatgen()
+
         composition = composition.remove_charges()
 
         elsyms = sorted({e.symbol for e in composition.elements})
@@ -114,8 +122,8 @@ def from_composition(
     @classmethod
     def from_structure(
         cls: Type[T],
-        meta_structure: Structure,
-        fields: Optional[List[str]] = None,
+        meta_structure: Structure | StructureReplica,
+        fields: Optional[list[str]] = None,
         **kwargs,
     ) -> T:
         fields = (
@@ -137,6 +145,11 @@ def from_structure(
             else fields
         )
         comp = meta_structure.composition.remove_charges()
+        if isinstance(comp, Composition) and not cls._use_pymatgen_rep:
+            comp = CompositionReplica.from_pymatgen(comp)
+        elif isinstance(comp, CompositionReplica) and cls._use_pymatgen_rep:
+            comp = comp.to_pymatgen()
+
         elsyms = sorted({e.symbol for e in comp.elements})
         symmetry = SymmetryData.from_structure(meta_structure)
 
@@ -158,6 +171,20 @@ def from_structure(
         return cls(**kwargs)
 
 
+class StructureMetadata(DbStructureMetadata):
+    """Pymatgen-object version of DbStructureMetadata."""
+
+    composition: Optional[Composition] = Field(
+        None, description="Full composition for the material."
+    )
+    composition_reduced: Optional[Composition] = Field(
+        None,
+        title="Reduced Composition",
+        description="Simplified representation of the composition.",
+    )
+    _use_pymatgen_rep: bool = True
+
+
 class MoleculeMetadata(EmmetBaseModel):
     """Mix-in class for molecule metadata."""
 
@@ -168,7 +195,7 @@ class MoleculeMetadata(EmmetBaseModel):
     natoms: Optional[int] = Field(
         None, description="Total number of atoms in the molecule"
     )
-    elements: Optional[List[Element]] = Field(
+    elements: Optional[list[Element]] = Field(
         None, description="List of elements in the molecule"
     )
     nelements: Optional[int] = Field(None, title="Number of Elements")
@@ -223,7 +250,7 @@ class MoleculeMetadata(EmmetBaseModel):
     def from_composition(
         cls: Type[S],
         comp: Composition,
-        fields: Optional[List[str]] = None,
+        fields: Optional[list[str]] = None,
         **kwargs,
     ) -> S:
         """
@@ -276,7 +303,7 @@ def from_composition(
     def from_molecule(
         cls: Type[S],
         meta_molecule: Molecule,
-        fields: Optional[List[str]] = None,
+        fields: Optional[list[str]] = None,
         **kwargs,
     ) -> S:
         fields = (