emmet-builders 0.84.7rc4__py3-none-any.whl → 0.84.9__py3-none-any.whl

emmet/builders/molecules/atomic.py

@@ -1,24 +1,29 @@
+from __future__ import annotations
+
 from collections import defaultdict
 from datetime import datetime
 from itertools import chain
 from math import ceil
-from typing import Optional, Iterable, Iterator, List, Dict
 
 from maggma.builders import Builder
 from maggma.core import Store
 from maggma.utils import grouper
 
-from emmet.core.qchem.task import TaskDocument
-from emmet.core.qchem.molecule import MoleculeDoc, evaluate_lot
+from emmet.builders.settings import EmmetBuildSettings
 from emmet.core.molecules.atomic import (
-    PartialChargesDoc,
-    PartialSpinsDoc,
     CHARGES_METHODS,
     SPINS_METHODS,
+    PartialChargesDoc,
+    PartialSpinsDoc,
 )
+from emmet.core.qchem.molecule import MoleculeDoc, evaluate_lot
+from emmet.core.qchem.task import TaskDocument
 from emmet.core.utils import jsanitize
-from emmet.builders.settings import EmmetBuildSettings
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator
 
 __author__ = "Evan Spotte-Smith"
 
@@ -55,9 +60,9 @@ class PartialChargesBuilder(Builder):
         tasks: Store,
         molecules: Store,
         charges: Store,
-        query: Optional[Dict] = None,
-        methods: Optional[List] = None,
-        settings: Optional[EmmetBuildSettings] = None,
+        query: dict | None = None,
+        methods: list | None = None,
+        settings: EmmetBuildSettings | None = None,
         **kwargs,
     ):
         self.tasks = tasks
@@ -105,7 +110,7 @@ class PartialChargesBuilder(Builder):
         self.charges.ensure_index("last_updated")
         self.charges.ensure_index("formula_alphabetical")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         """Prechunk the builder for distributed computation"""
 
         temp_query = dict(self.query)
@@ -131,7 +136,7 @@ class PartialChargesBuilder(Builder):
             query["species_hash"] = {"$in": list(hash_chunk)}
             yield {"query": query}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets all items to process into partial charges documents.
         This does no datetime checking; relying on on whether
@@ -178,12 +183,12 @@ class PartialChargesBuilder(Builder):
 
             yield molecules
 
-    def process_item(self, items: List[Dict]) -> List[Dict]:
+    def process_item(self, items: list[dict]) -> list[dict]:
         """
         Process the tasks into PartialChargesDocs
 
         Args:
-            tasks List[Dict] : a list of MoleculeDocs in dict form
+            tasks list[dict] : a list of MoleculeDocs in dict form
 
         Returns:
             [dict] : a list of new partial charges docs
@@ -275,7 +280,7 @@ class PartialChargesBuilder(Builder):
 
         return jsanitize([doc.model_dump() for doc in charges_docs], allow_bson=True)
 
-    def update_targets(self, items: List[List[Dict]]):
+    def update_targets(self, items: list[list[dict]]):
         """
         Inserts the new documents into the charges collection
 
@@ -334,9 +339,9 @@ class PartialSpinsBuilder(Builder):
         tasks: Store,
         molecules: Store,
         spins: Store,
-        query: Optional[Dict] = None,
-        methods: Optional[List] = None,
-        settings: Optional[EmmetBuildSettings] = None,
+        query: dict | None = None,
+        methods: list | None = None,
+        settings: EmmetBuildSettings | None = None,
         **kwargs,
     ):
         self.tasks = tasks
@@ -384,7 +389,7 @@ class PartialSpinsBuilder(Builder):
         self.spins.ensure_index("last_updated")
         self.spins.ensure_index("formula_alphabetical")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         """Prechunk the builder for distributed computation"""
 
         temp_query = dict(self.query)
@@ -410,7 +415,7 @@ class PartialSpinsBuilder(Builder):
             query["species_hash"] = {"$in": list(hash_chunk)}
             yield {"query": query}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets all items to process into partial spins documents.
         This does no datetime checking; relying on on whether
@@ -457,12 +462,12 @@ class PartialSpinsBuilder(Builder):
 
             yield molecules
 
-    def process_item(self, items: List[Dict]) -> List[Dict]:
+    def process_item(self, items: list[dict]) -> list[dict]:
         """
         Process the tasks into PartialSpinsDocs
 
         Args:
-            tasks List[Dict] : a list of MoleculeDocs in dict form
+            tasks list[dict] : a list of MoleculeDocs in dict form
 
         Returns:
             [dict] : a list of new partial spins docs
@@ -555,7 +560,7 @@ class PartialSpinsBuilder(Builder):
 
         return jsanitize([doc.model_dump() for doc in spins_docs], allow_bson=True)
 
-    def update_targets(self, items: List[List[Dict]]):
+    def update_targets(self, items: list[list[dict]]):
         """
         Inserts the new documents into the spins collection
 

emmet/builders/molecules/bonds.py

@@ -1,19 +1,24 @@
+from __future__ import annotations
+
 from collections import defaultdict
 from datetime import datetime
 from itertools import chain
 from math import ceil
-from typing import Optional, Iterable, Iterator, List, Dict
 
 from maggma.builders import Builder
 from maggma.core import Store
 from maggma.utils import grouper
 
-from emmet.core.qchem.task import TaskDocument
+from emmet.builders.settings import EmmetBuildSettings
+from emmet.core.molecules.bonds import BOND_METHODS, MoleculeBondingDoc
 from emmet.core.qchem.molecule import MoleculeDoc, evaluate_lot
-from emmet.core.molecules.bonds import MoleculeBondingDoc, BOND_METHODS
+from emmet.core.qchem.task import TaskDocument
 from emmet.core.utils import jsanitize
-from emmet.builders.settings import EmmetBuildSettings
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator
 
 __author__ = "Evan Spotte-Smith"
 
@@ -54,9 +59,9 @@ class BondingBuilder(Builder):
         tasks: Store,
         molecules: Store,
         bonds: Store,
-        query: Optional[Dict] = None,
-        methods: Optional[List] = None,
-        settings: Optional[EmmetBuildSettings] = None,
+        query: dict | None = None,
+        methods: list | None = None,
+        settings: EmmetBuildSettings | None = None,
         **kwargs,
     ):
         self.tasks = tasks
@@ -104,7 +109,7 @@ class BondingBuilder(Builder):
         self.bonds.ensure_index("last_updated")
         self.bonds.ensure_index("formula_alphabetical")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         """Prechunk the builder for distributed computation"""
 
         temp_query = dict(self.query)
@@ -130,7 +135,7 @@ class BondingBuilder(Builder):
             query["species_hash"] = {"$in": list(hash_chunk)}
             yield {"query": query}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets all items to process into bonding documents.
         This does no datetime checking; relying on on whether
@@ -177,12 +182,12 @@ class BondingBuilder(Builder):
 
             yield molecules
 
-    def process_item(self, items: List[Dict]) -> List[Dict]:
+    def process_item(self, items: list[dict]) -> list[dict]:
         """
         Process the tasks into MoleculeBondingDocs
 
         Args:
-            tasks List[Dict] : a list of MoleculeDocs in dict form
+            tasks list[dict] : a list of MoleculeDocs in dict form
 
         Returns:
             [dict] : a list of new bonding docs
@@ -292,7 +297,7 @@ class BondingBuilder(Builder):
 
         return jsanitize([doc.model_dump() for doc in bonding_docs], allow_bson=True)
 
-    def update_targets(self, items: List[List[Dict]]):
+    def update_targets(self, items: list[list[dict]]):
         """
         Inserts the new documents into the bonds collection
 

emmet/builders/molecules/electric.py

@@ -1,19 +1,24 @@
+from __future__ import annotations
+
 from collections import defaultdict
 from datetime import datetime
 from itertools import chain
 from math import ceil
-from typing import Optional, Iterable, Iterator, List, Dict
 
 from maggma.builders import Builder
 from maggma.core import Store
 from maggma.utils import grouper
 
-from emmet.core.qchem.task import TaskDocument
-from emmet.core.qchem.molecule import MoleculeDoc, evaluate_lot
+from emmet.builders.settings import EmmetBuildSettings
 from emmet.core.molecules.electric import ElectricMultipoleDoc
+from emmet.core.qchem.molecule import MoleculeDoc, evaluate_lot
+from emmet.core.qchem.task import TaskDocument
 from emmet.core.utils import jsanitize
-from emmet.builders.settings import EmmetBuildSettings
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator
 
 __author__ = "Evan Spotte-Smith"
 
@@ -41,8 +46,8 @@ class ElectricMultipoleBuilder(Builder):
         tasks: Store,
         molecules: Store,
         multipoles: Store,
-        query: Optional[Dict] = None,
-        settings: Optional[EmmetBuildSettings] = None,
+        query: dict | None = None,
+        settings: EmmetBuildSettings | None = None,
         **kwargs,
     ):
         self.tasks = tasks
@@ -89,7 +94,7 @@ class ElectricMultipoleBuilder(Builder):
         self.multipoles.ensure_index("last_updated")
         self.multipoles.ensure_index("formula_alphabetical")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         """Prechunk the builder for distributed computation"""
 
         temp_query = dict(self.query)
@@ -115,7 +120,7 @@ class ElectricMultipoleBuilder(Builder):
             query["species_hash"] = {"$in": list(hash_chunk)}
             yield {"query": query}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets all items to process into multipole documents.
         This does no datetime checking; relying on on whether
@@ -162,12 +167,12 @@ class ElectricMultipoleBuilder(Builder):
 
             yield molecules
 
-    def process_item(self, items: List[Dict]) -> List[Dict]:
+    def process_item(self, items: list[dict]) -> list[dict]:
         """
         Process the tasks into ElectricMultipoleDocs
 
         Args:
-            tasks List[Dict] : a list of MoleculeDocs in dict form
+            tasks list[dict] : a list of MoleculeDocs in dict form
 
         Returns:
             [dict] : a list of new electric multipole docs
@@ -250,7 +255,7 @@ class ElectricMultipoleBuilder(Builder):
 
         return jsanitize([doc.model_dump() for doc in multipole_docs], allow_bson=True)
 
-    def update_targets(self, items: List[List[Dict]]):
+    def update_targets(self, items: list[list[dict]]):
         """
         Inserts the new documents into the multipoles collection
 

emmet/builders/molecules/metal_binding.py

@@ -1,24 +1,27 @@
+from __future__ import annotations
+
+import copy
 from datetime import datetime
 from itertools import chain
 from math import ceil
-from typing import Optional, Iterable, Iterator, List, Dict
-import copy
-
-from pymatgen.core.structure import Molecule
-from pymatgen.util.graph_hashing import weisfeiler_lehman_graph_hash
+from typing import TYPE_CHECKING
 
 from maggma.builders import Builder
 from maggma.core import Store
 from maggma.utils import grouper
+from pymatgen.core.structure import Molecule
+from pymatgen.util.graph_hashing import weisfeiler_lehman_graph_hash
 
-from emmet.core.qchem.molecule import MoleculeDoc
+from emmet.builders.settings import EmmetBuildSettings
 from emmet.core.molecules.atomic import PartialChargesDoc, PartialSpinsDoc
 from emmet.core.molecules.bonds import MoleculeBondingDoc, metals
+from emmet.core.molecules.metal_binding import METAL_BINDING_METHODS, MetalBindingDoc
 from emmet.core.molecules.thermo import MoleculeThermoDoc
-from emmet.core.molecules.metal_binding import MetalBindingDoc, METAL_BINDING_METHODS
+from emmet.core.qchem.molecule import MoleculeDoc
 from emmet.core.utils import jsanitize
-from emmet.builders.settings import EmmetBuildSettings
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator
 
 __author__ = "Evan Spotte-Smith"
 
@@ -73,9 +76,9 @@ class MetalBindingBuilder(Builder):
         bonds: Store,
         thermo: Store,
         metal_binding: Store,
-        query: Optional[Dict] = None,
-        methods: Optional[List] = None,
-        settings: Optional[EmmetBuildSettings] = None,
+        query: dict | None = None,
+        methods: list | None = None,
+        settings: EmmetBuildSettings | None = None,
         **kwargs,
     ):
         self.molecules = molecules
@@ -161,7 +164,7 @@ class MetalBindingBuilder(Builder):
         self.metal_binding.ensure_index("formula_alphabetical")
         self.metal_binding.ensure_index("method")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         """Prechunk the builder for distributed computation"""
 
         temp_query = dict(self.query)
@@ -187,7 +190,7 @@ class MetalBindingBuilder(Builder):
             query["species_hash"] = {"$in": list(hash_chunk)}
             yield {"query": query}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets all items to process into metal_binding documents.
 
@@ -232,12 +235,12 @@ class MetalBindingBuilder(Builder):
 
             yield molecules
 
-    def process_item(self, items: List[Dict]) -> List[Dict]:
+    def process_item(self, items: list[dict]) -> list[dict]:
         """
         Process molecule, bonding, partial charges, partial spins, and thermo documents into MetalBindingDocs
 
         Args:
-            tasks List[Dict] : a list of MoleculeDocs in dict form
+            tasks list[dict] : a list of MoleculeDocs in dict form
 
         Returns:
             [dict] : a list of new metal binding docs
@@ -491,7 +494,7 @@ class MetalBindingBuilder(Builder):
 
         return jsanitize([doc.model_dump() for doc in binding_docs], allow_bson=True)
 
-    def update_targets(self, items: List[List[Dict]]):
+    def update_targets(self, items: list[list[dict]]):
         """
         Inserts the new documents into the metal_binding collection
 

emmet/builders/molecules/orbitals.py

@@ -1,19 +1,23 @@
+from __future__ import annotations
+
 from collections import defaultdict
 from datetime import datetime
 from itertools import chain
 from math import ceil
-from typing import Optional, Iterable, Iterator, List, Dict
+from typing import TYPE_CHECKING
 
 from maggma.builders import Builder
 from maggma.core import Store
 from maggma.utils import grouper
 
-from emmet.core.qchem.task import TaskDocument
-from emmet.core.qchem.molecule import MoleculeDoc, evaluate_lot
+from emmet.builders.settings import EmmetBuildSettings
 from emmet.core.molecules.orbitals import OrbitalDoc
+from emmet.core.qchem.molecule import MoleculeDoc, evaluate_lot
+from emmet.core.qchem.task import TaskDocument
 from emmet.core.utils import jsanitize
-from emmet.builders.settings import EmmetBuildSettings
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator
 
 __author__ = "Evan Spotte-Smith"
 
@@ -40,8 +44,8 @@ class OrbitalBuilder(Builder):
         tasks: Store,
         molecules: Store,
         orbitals: Store,
-        query: Optional[Dict] = None,
-        settings: Optional[EmmetBuildSettings] = None,
+        query: dict | None = None,
+        settings: EmmetBuildSettings | None = None,
         **kwargs,
     ):
         self.tasks = tasks
@@ -87,7 +91,7 @@ class OrbitalBuilder(Builder):
         self.orbitals.ensure_index("last_updated")
         self.orbitals.ensure_index("formula_alphabetical")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         """Prechunk the builder for distributed computation"""
 
         temp_query = dict(self.query)
@@ -113,7 +117,7 @@ class OrbitalBuilder(Builder):
             query["species_hash"] = {"$in": list(hash_chunk)}
             yield {"query": query}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets all items to process into orbital documents.
         This does no datetime checking; relying on on whether
@@ -160,12 +164,12 @@ class OrbitalBuilder(Builder):
 
             yield molecules
 
-    def process_item(self, items: List[Dict]) -> List[Dict]:
+    def process_item(self, items: list[dict]) -> list[dict]:
         """
         Process the tasks into a OrbitalDocs
 
         Args:
-            tasks List[Dict] : a list of MoleculeDocs in dict form
+            tasks list[dict] : a list of MoleculeDocs in dict form
 
         Returns:
             [dict] : a list of new orbital docs
@@ -257,7 +261,7 @@ class OrbitalBuilder(Builder):
 
         return jsanitize([doc.model_dump() for doc in orbital_docs], allow_bson=True)
 
-    def update_targets(self, items: List[List[Dict]]):
+    def update_targets(self, items: list[list[dict]]):
         """
         Inserts the new documents into the orbitals collection
 

emmet/builders/molecules/redox.py

@@ -1,22 +1,28 @@
-from collections import defaultdict
+from __future__ import annotations
+
 import copy
+from collections import defaultdict
 from datetime import datetime
 from itertools import chain, groupby
 from math import ceil
-from typing import Any, Dict, Iterable, Iterator, List, Optional, Union
 
 from maggma.builders import Builder
 from maggma.core import Store
 from maggma.utils import grouper
 
-from emmet.core.qchem.task import TaskDocument
-from emmet.core.qchem.molecule import MoleculeDoc
+from emmet.builders.settings import EmmetBuildSettings
 from emmet.core.molecules.bonds import metals
-from emmet.core.molecules.thermo import MoleculeThermoDoc
 from emmet.core.molecules.redox import RedoxDoc
+from emmet.core.molecules.thermo import MoleculeThermoDoc
+from emmet.core.qchem.molecule import MoleculeDoc
+from emmet.core.qchem.task import TaskDocument
 from emmet.core.utils import confirm_molecule, get_graph_hash, jsanitize
-from emmet.builders.settings import EmmetBuildSettings
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator
+    from typing import Any
 
 __author__ = "Evan Spotte-Smith"
 
@@ -51,8 +57,8 @@ class RedoxBuilder(Builder):
         molecules: Store,
         thermo: Store,
         redox: Store,
-        query: Optional[Dict] = None,
-        settings: Optional[EmmetBuildSettings] = None,
+        query: dict | None = None,
+        settings: EmmetBuildSettings | None = None,
         **kwargs,
     ):
         self.tasks = tasks
@@ -107,7 +113,7 @@ class RedoxBuilder(Builder):
         self.redox.ensure_index("last_updated")
         self.redox.ensure_index("formula_alphabetical")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         """Prechunk the builder for distributed computation"""
 
         temp_query = dict(self.query)
@@ -133,7 +139,7 @@ class RedoxBuilder(Builder):
             query["species_hash"] = {"$in": list(hash_chunk)}
             yield {"query": query}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets all items to process into redox documents.
         This does no datetime checking; relying on on whether
@@ -180,12 +186,12 @@ class RedoxBuilder(Builder):
 
             yield molecules
 
-    def process_item(self, items: List[Dict]) -> List[Dict]:
+    def process_item(self, items: list[dict]) -> list[dict]:
         """
         Process the tasks into a RedoxDoc
 
         Args:
-            tasks List[Dict] : a list of MoleculeDocs in dict form
+            tasks list[dict] : a list of MoleculeDocs in dict form
 
         Returns:
             [dict] : a list of new redox docs
@@ -203,7 +209,7 @@ class RedoxBuilder(Builder):
 
         for graph_group in group_by_graph.values():
             # Molecule docs will be grouped by charge
-            charges: Dict[int, Any] = dict()
+            charges: dict[int, Any] = dict()
 
             for gg in graph_group:
                 # First, grab relevant MoleculeThermoDocs and identify possible IE/EA single-points
@@ -360,7 +366,7 @@ class RedoxBuilder(Builder):
             [doc.model_dump() for doc in redox_docs if doc is not None], allow_bson=True
         )
 
-    def update_targets(self, items: List[List[Dict]]):
+    def update_targets(self, items: list[list[dict]]):
         """
         Inserts the new documents into the orbitals collection
 
@@ -391,7 +397,7 @@ class RedoxBuilder(Builder):
             self.logger.info("No items to update")
 
     @staticmethod
-    def _group_by_graph(mol_docs: List[MoleculeDoc]) -> Dict[int, List[MoleculeDoc]]:
+    def _group_by_graph(mol_docs: list[MoleculeDoc]) -> dict[int, list[MoleculeDoc]]:
         """
         Group molecule docs by molecular graph connectivity
 
@@ -399,7 +405,7 @@ class RedoxBuilder(Builder):
         :return: Grouped molecule entries
         """
 
-        graph_hashes_nometal: List[str] = list()
+        graph_hashes_nometal: list[str] = list()
         results = defaultdict(list)
 
         # Within each group, group by the covalent molecular graph
@@ -430,10 +436,10 @@ class RedoxBuilder(Builder):
 
     @staticmethod
     def _collect_by_lot_solvent(
-        thermo_docs: List[MoleculeThermoDoc],
-        ie_docs: List[TaskDocument],
-        ea_docs: List[TaskDocument],
-    ) -> Dict[str, Any]:
+        thermo_docs: list[MoleculeThermoDoc],
+        ie_docs: list[TaskDocument],
+        ea_docs: list[TaskDocument],
+    ) -> dict[str, Any]:
         """
         For a given MoleculeDoc, group potential MoleculeThermoDocs and TaskDocs for
         IE/EA calculations based on level of theory and solvent.
@@ -452,7 +458,7 @@ class RedoxBuilder(Builder):
                  }
         """
 
-        def _lot_solv(doc: Union[MoleculeThermoDoc, TaskDocument]):
+        def _lot_solv(doc: MoleculeThermoDoc | TaskDocument):
             if isinstance(doc, MoleculeThermoDoc):
                 if doc.correction:
                     return doc.correction_lot_solvent