From c4d0c15fc47e773d2a093953c5df940383f26b63 Mon Sep 17 00:00:00 2001 From: Shyam D Date: Fri, 19 Feb 2021 14:28:33 -0800 Subject: [PATCH] fix composition validators --- emmet-core/emmet/core/stubs.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/emmet-core/emmet/core/stubs.py b/emmet-core/emmet/core/stubs.py index 79179f2155..bfa8d3326b 100644 --- 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