chore: types and pydantic

Exabyte-io · Sep 24, 2024 · d36a741 · d36a741
1 parent 6b651e1
commit d36a741
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 11 deletions.
diff --git a/src/py/mat3ra/made/tools/build/__init__.py b/src/py/mat3ra/made/tools/build/__init__.py
@@ -27,6 +27,11 @@ class BaseSelectorParameters(BaseModel):
     default_index: int = 0
 
 
+class BaseBuilderParameters(BaseModel):
+    class Config:
+        arbitrary_types_allowed = True
+
+
 class BaseBuilder(BaseModel):
     """
     Base class for material builders.

diff --git a/src/py/mat3ra/made/tools/build/grain_boundary/builders.py b/src/py/mat3ra/made/tools/build/grain_boundary/builders.py
@@ -24,7 +24,7 @@ class GrainBoundaryBuilder(ZSLStrainMatchingInterfaceBuilder):
     3. Finally, creating a slab from the rotated interface.
     """
 
-    _BuildParametersType = GrainBoundaryBuilderParameters
+    _BuildParametersType: type(GrainBoundaryBuilderParameters) = GrainBoundaryBuilderParameters  # type: ignore
     _ConfigurationType: type(GrainBoundaryConfiguration) = GrainBoundaryConfiguration  # type: ignore
     _GeneratedItemType: type(Material) = Material  # type: ignore
     selector_parameters: type(  # type: ignore

diff --git a/src/py/mat3ra/made/tools/build/grain_boundary/configuration.py b/src/py/mat3ra/made/tools/build/grain_boundary/configuration.py
@@ -1,11 +1,11 @@
-from pydantic import BaseModel
-from mat3ra.code.entity import InMemoryEntity
+from typing import Optional
 
+from .. import BaseConfiguration
 from ..slab.configuration import SlabConfiguration
 from ..slab.termination import Termination
 
 
-class GrainBoundaryConfiguration(BaseModel, InMemoryEntity):
+class GrainBoundaryConfiguration(BaseConfiguration):
     """
     Configuration for a grain boundary in a slab material.
 
@@ -26,7 +26,7 @@ class GrainBoundaryConfiguration(BaseModel, InMemoryEntity):
     phase_2_termination: Termination
     gap: float = 3.0
     slab_configuration: SlabConfiguration
-    slab_termination: Termination
+    slab_termination: Optional[Termination]
 
     @property
     def _json(self):

diff --git a/src/py/mat3ra/made/tools/build/interface/builders.py b/src/py/mat3ra/made/tools/build/interface/builders.py
@@ -22,7 +22,7 @@
 from ..slab.configuration import SlabConfiguration
 from ...analyze import get_chemical_formula
 from ...convert import to_ase, from_ase, to_pymatgen, PymatgenInterface, ASEAtoms
-from ...build import BaseBuilder
+from ...build import BaseBuilder, BaseBuilderParameters
 
 
 class InterfaceBuilderParameters(BaseModel):
@@ -106,7 +106,7 @@ def _post_process(self, items: List[_GeneratedItemType], post_process_parameters
 ########################################################################################
 #                       Strain Matching Interface Builders                             #
 ########################################################################################
-class StrainMatchingInterfaceBuilderParameters(BaseModel):
+class StrainMatchingInterfaceBuilderParameters(BaseBuilderParameters):
     strain_matching_parameters: Optional[Any] = None
 
 

diff --git a/src/py/mat3ra/made/tools/build/slab/configuration.py b/src/py/mat3ra/made/tools/build/slab/configuration.py
@@ -1,16 +1,15 @@
 from typing import List, Tuple, Any
 
 import numpy as np
-from pydantic import BaseModel
-
-from mat3ra.code.entity import InMemoryEntity
 
 from mat3ra.made.material import Material
+
+from .. import BaseConfiguration
 from ...third_party import PymatgenSpacegroupAnalyzer
 from ...convert import to_pymatgen, from_pymatgen
 
 
-class SlabConfiguration(BaseModel, InMemoryEntity):
+class SlabConfiguration(BaseConfiguration):
     """
     Configuration for building a slab.