[WIP] Update static collections

emmet-core/emmet/core/grain_boundary.py

@@ -1,9 +1,18 @@
+from __future__ import annotations
+
 from pydantic import BaseModel, Field
+from pymatgen.core import Structure
+from typing import TYPE_CHECKING
 
 from emmet.core.types.enums import ValueEnum
 from emmet.core.types.pymatgen_types.grain_boundary_adapter import GrainBoundaryType
+from emmet.core.types.pymatgen_types.structure_adapter import StructureType
 from emmet.core.types.typing import DateTimeType
 
+if TYPE_CHECKING:
+    from typing import Any
+    from typing_extensions import Self
+
 
 class GBTypeEnum(ValueEnum):
     """
@@ -68,7 +77,7 @@ class GrainBoundaryDoc(BaseModel):
 
     w_sep: float | None = Field(None, description="Work of separation in J/m^2.")
 
-    cif: str | None = Field(None, description="CIF file of the structure.")
+    structure: StructureType | None = Field(None, description="Structure.")
 
     chemsys: str | None = Field(
         None, description="Dash-delimited string of elements in the material."
@@ -77,3 +86,19 @@ class GrainBoundaryDoc(BaseModel):
     last_updated: DateTimeType = Field(
         description="Timestamp for the most recent calculation for this Material document.",
     )
+
+    @classmethod
+    def _migrate_schema(cls, config: dict[str, Any]) -> Self:
+        """Pipe legacy CIF data into a pymatgen structure."""
+        if isinstance(cif_str := config.pop("cif", None), str) and not config.get(
+            "structure"
+        ):
+            config["structure"] = Structure.from_str(cif_str, fmt="cif")
+        return cls(**config)
+
+    @property
+    def cif(self) -> str | None:
+        """Support accessing legacy CIF from structure field."""
+        if self.structure:
+            return self.structure.to(fmt="cif")
+        return None

emmet-core/emmet/core/surface_properties.py

@@ -1,7 +1,17 @@
-from pydantic import BaseModel, Field
+"""Define schemas for surfaces and heterostructures."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field, field_validator
+from typing import TYPE_CHECKING
+
+from pymatgen.core import Structure
 
 from emmet.core.types.pymatgen_types.structure_adapter import StructureType
 
+if TYPE_CHECKING:
+    from typing import Any
+
 
 class SurfaceEntry(BaseModel):
     """
@@ -28,9 +38,9 @@ class SurfaceEntry(BaseModel):
         description="Whether it is a reconstructed surface.",
     )
 
-    structure: str | None = Field(
+    structure: StructureType | None = Field(
         None,
-        description="CIF of slab structure.",
+        description="Slab structure.",
     )
 
     work_function: float | None = Field(
@@ -53,6 +63,13 @@ class SurfaceEntry(BaseModel):
         description="Whether the surface has wulff entry.",
     )
 
+    @field_validator("structure", mode="before")
+    def get_structure_from_cif(cls, v: Any) -> Structure | None:
+        """Transform legacy CIF data to pymatgen Structure."""
+        if isinstance(v, str):
+            return Structure.from_str(v, fmt="cif")
+        return v
+
 
 class SurfacePropDoc(BaseModel):
     """

emmet-core/emmet/core/synthesis/core.py

@@ -1,12 +1,9 @@
-from typing import Any
-
 from pydantic import BaseModel, Field
 
 from emmet.core.synthesis.materials import ExtractedMaterial
 from emmet.core.synthesis.operations import Operation
 from emmet.core.synthesis.reaction import ReactionFormula
 from emmet.core.types.enums import ValueEnum
-from emmet.core.utils import arrow_incompatible
 
 
 class SynthesisTypeEnum(str, ValueEnum):
@@ -60,7 +57,21 @@ class SynthesisRecipe(BaseModel):
     )
 
 
-@arrow_incompatible
+class AtlasSearchText(BaseModel):
+    """Schematize MongoDB Atlas search text match."""
+
+    value: str | None = None
+    type: str | None = None
+
+
+class AtlasSearchHighlight(BaseModel):
+    """Schematize MongoDB Atlas search highlight match."""
+
+    score: float | None = None
+    path: str | None = None
+    texts: list[AtlasSearchText] | None = None
+
+
 class SynthesisSearchResultModel(SynthesisRecipe):
     """
     Model for a document containing synthesis recipes
@@ -71,7 +82,7 @@ class SynthesisSearchResultModel(SynthesisRecipe):
         None,
         description="Search score.",
     )
-    highlights: list[Any] | None = Field(
+    highlights: list[AtlasSearchHighlight] | None = Field(
         None,
         description="Search highlights.",
     )

emmet-core/tests/test_model_fields.py

@@ -259,7 +259,7 @@
         "final_structure",
         "pretty_formula",
         "w_sep",
-        "cif",
+        "structure",
         "chemsys",
         "last_updated",
     ],