-
Notifications
You must be signed in to change notification settings - Fork 62
/
__init__.py
37 lines (31 loc) · 1.52 KB
/
__init__.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 in pydantic models for common MSONable classes, particularly those in Pymatgen
Use pymatgen classes in pydantic models by importing them from there when you need schema
"""
from pymatgen.analysis.diffraction.xrd import DiffractionPattern
from pymatgen.analysis.xas.spectrum import XAS
from pymatgen.core.structure import Composition, Lattice, Structure
from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
from emmet.stubs.utils import patch_msonable, use_model
from emmet.stubs.math import Matrix3D, Vector3D
from emmet.stubs.structure import Lattice as StubLattice
from emmet.stubs.structure import Structure as StubStructure
from emmet.stubs.entries import Composition as StubComposition
from emmet.stubs.entries import ComputedEntry as StubComputedEntry
from emmet.stubs.entries import ComputedStructureEntry as StubComputedStructureEntry
from emmet.stubs.xrd import XRDPattern as StubXRDPattern
"""
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
"""
use_model(Structure, StubStructure)
use_model(Lattice, StubLattice)
use_model(Composition, StubComposition, add_monty=False)
use_model(ComputedEntry, StubComputedEntry)
use_model(DiffractionPattern, StubXRDPattern)
use_model(ComputedStructureEntry, StubComputedStructureEntry)
# This is after the main block since it depends on that
from emmet.stubs.xas import XASSpectrum # noqa
use_model(XAS, XASSpectrum)