fix composition validators

materialsproject · Feb 19, 2021 · c4d0c15 · c4d0c15
1 parent 7a4f692
commit c4d0c15
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/emmet-core/emmet/core/stubs.py b/emmet-core/emmet/core/stubs.py
@@ -21,4 +21,16 @@ class StubComposition(BaseModel):
     __root__: Dict[Element, float]
 
 
+@classmethod
+def get_validators(cls):
+    yield validate_composition
+
+
+def validate_composition(cls, v):
+    if isinstance(v, pymatgen.core.structure.Composition):
+        return v
+    return pymatgen.core.structure.Composition(**v)
+
+
 pymatgen.core.structure.Composition.__pydantic_model__ = StubComposition
+pymatgen.core.structure.Composition.__get_validators__ = get_validators