diff --git a/emmet-core/emmet/core/grain_boundary.py b/emmet-core/emmet/core/grain_boundary.py index 1b8779471c..208636a4ab 100644 --- a/emmet-core/emmet/core/grain_boundary.py +++ b/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 diff --git a/emmet-core/emmet/core/surface_properties.py b/emmet-core/emmet/core/surface_properties.py index 3d82f5569b..3adb420e76 100644 --- a/emmet-core/emmet/core/surface_properties.py +++ b/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): """ diff --git a/emmet-core/emmet/core/synthesis/core.py b/emmet-core/emmet/core/synthesis/core.py index 1c3f605526..16d5d0e391 100644 --- a/emmet-core/emmet/core/synthesis/core.py +++ b/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.", ) diff --git a/emmet-core/tests/test_model_fields.py b/emmet-core/tests/test_model_fields.py index eeae509150..de15b9495f 100644 --- a/emmet-core/tests/test_model_fields.py +++ b/emmet-core/tests/test_model_fields.py @@ -259,7 +259,7 @@ "final_structure", "pretty_formula", "w_sep", - "cif", + "structure", "chemsys", "last_updated", ],