-
Notifications
You must be signed in to change notification settings - Fork 62
/
stubs.py
37 lines (26 loc) · 972 Bytes
/
stubs.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
# isort: off
"""
This module stubs some pymatgen classes that implement custom behavior
outside the standard MSONable model
"""
from typing import Dict
import pymatgen.core.structure
from pydantic import BaseModel
from pymatgen.core.periodic_table import Element
"""
The stub names are kept in sync with the actual classes so they
show up correctly in the JSON Schema. They are imported here
in as Stubbed classes to prevent name clashing
"""
class StubComposition(BaseModel):
"""A dictionary mapping element to total quantity"""
__root__: Dict[Element, float]
@classmethod # type: ignore
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