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

emmet/builders/abinit/phonon.py

@@ -1,7 +1,6 @@
 import os
 import tempfile
 from math import ceil
-from typing import Dict, Iterator, List, Optional, Tuple
 
 import numpy as np
 from abipy.abio.inputs import AnaddbInput
@@ -36,6 +35,11 @@ from emmet.core.phonon import (
 from emmet.core.polar import BornEffectiveCharges, DielectricDoc, IRDielectric
 from emmet.core.utils import jsanitize
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 SETTINGS = EmmetBuildSettings()
 
 
@@ -50,8 +54,8 @@ class PhononBuilder(Builder):
         ddb_files: Store,
         th_disp: Store,
         phonon_website: Store,
-        query: Optional[Dict] = None,
-        manager: Optional[TaskManager] = None,
+        query: dict | None = None,
+        manager: TaskManager | None = None,
         symprec: float = SETTINGS.SYMPREC,
         angle_tolerance: float = SETTINGS.ANGLE_TOL,
         chunk_size=100,
@@ -134,7 +138,7 @@ class PhononBuilder(Builder):
         for mpid_chunk in grouper(mats, N):
             yield {"query": {self.phonon_materials.key: {"$in": list(mpid_chunk)}}}
 
-    def get_items(self) -> Iterator[Dict]:
+    def get_items(self) -> Iterator[dict]:
         """
         Gets all materials that need phonons
 
@@ -188,7 +192,7 @@ class PhononBuilder(Builder):
 
             yield item
 
-    def process_item(self, item: Dict) -> Optional[Dict]:
+    def process_item(self, item: dict) -> dict | None:
         """
         Generates the full phonon document from an item
 
@@ -293,7 +297,7 @@ class PhononBuilder(Builder):
             )
             return None
 
-    def get_phonon_properties(self, item: Dict) -> Dict:
+    def get_phonon_properties(self, item: dict) -> dict:
         """
         Extracts the phonon properties from the item
         """
@@ -511,7 +515,7 @@ class PhononBuilder(Builder):
         dos: str = "tetra",
         lo_to_splitting: bool = True,
         use_dieflag: bool = True,
-    ) -> Tuple[AnaddbInput, Optional[List]]:
+    ) -> tuple[AnaddbInput, list | None]:
         """
         creates the AnaddbInput object to calculate the phonon properties.
         It also returns the list of qpoints labels for generating the PhononBandStructureSymmLine.
@@ -637,7 +641,7 @@ class PhononBuilder(Builder):
 
     @staticmethod
     def get_pmg_bs(
-        phbands: PhononBands, labels_list: List
+        phbands: PhononBands, labels_list: list
     ) -> PhononBandStructureSymmLine:
         """
         Generates a PhononBandStructureSymmLine starting from a abipy PhononBands object
@@ -726,7 +730,7 @@ class PhononBuilder(Builder):
 
         return data
 
-    def update_targets(self, items: List[Dict]):
+    def update_targets(self, items: list[dict]):
         """
         Inserts the new task_types into the task_types collection
 
@@ -769,7 +773,7 @@ class PhononBuilder(Builder):
 
 def get_warnings(
     asr_break: float, cnsr_break: float, ph_bs: PhononBandStructureSymmLine
-) -> List[PhononWarnings]:
+) -> list[PhononWarnings]:
     """
 
     Args:
@@ -810,7 +814,7 @@ def get_warnings(
 
 def get_thermodynamic_properties(
     ph_dos: CompletePhononDos,
-) -> Tuple[ThermodynamicProperties, VibrationalEnergy]:
+) -> tuple[ThermodynamicProperties, VibrationalEnergy]:
     """
     Calculates the thermodynamic properties from a phonon DOS
 

emmet/builders/abinit/sound_velocity.py

@@ -1,18 +1,22 @@
 import tempfile
 import traceback
 from math import ceil
-from maggma.utils import grouper
-from typing import Optional, Dict, List, Iterator
 
-from abipy.dfpt.vsound import SoundVelocity as AbiSoundVelocity
 from abipy.dfpt.ddb import DdbFile
+from abipy.dfpt.vsound import SoundVelocity as AbiSoundVelocity
+from abipy.flowtk.tasks import TaskManager
 from maggma.builders import Builder
 from maggma.core import Store
-from abipy.flowtk.tasks import TaskManager
+from maggma.utils import grouper
 
 from emmet.core.phonon import SoundVelocity
 from emmet.core.utils import jsanitize
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 
 class SoundVelocityBuilder(Builder):
     def __init__(
@@ -20,8 +24,8 @@ class SoundVelocityBuilder(Builder):
         phonon_materials: Store,
         ddb_source: Store,
         sound_vel: Store,
-        query: Optional[dict] = None,
-        manager: Optional[TaskManager] = None,
+        query: dict | None = None,
+        manager: TaskManager | None = None,
         **kwargs
     ):
         """
@@ -72,7 +76,7 @@ class SoundVelocityBuilder(Builder):
         for mpid_chunk in grouper(mats, N):
             yield {"query": {self.phonon_materials.key: {"$in": list(mpid_chunk)}}}
 
-    def get_items(self) -> Iterator[Dict]:
+    def get_items(self) -> Iterator[dict]:
         """
         Gets all materials that need sound velocity.
 
@@ -117,7 +121,7 @@ class SoundVelocityBuilder(Builder):
 
             yield item
 
-    def process_item(self, item: Dict) -> Optional[Dict]:
+    def process_item(self, item: dict) -> dict | None:
         """
         Generates the sound velocity document from an item
 
@@ -153,7 +157,7 @@ class SoundVelocityBuilder(Builder):
             return None
 
     @staticmethod
-    def get_sound_vel(item: Dict) -> Dict:
+    def get_sound_vel(item: dict) -> dict:
         """
         Runs anaddb and return the extracted data for the speed of sound.
 
@@ -187,7 +191,7 @@ class SoundVelocityBuilder(Builder):
 
             return sv_data
 
-    def update_targets(self, items: List[Dict]):
+    def update_targets(self, items: list[dict]):
         """
         Inserts the new task_types into the task_types collection
 

emmet/builders/feff/xas.py

@@ -1,7 +1,6 @@
 import traceback
 from datetime import datetime
 from itertools import chain
-from typing import Dict, List
 
 from maggma.builders import GroupBuilder
 from maggma.core import Store
@@ -27,7 +26,7 @@ class XASBuilder(GroupBuilder):
         super().__init__(source=tasks, target=xas, grouping_keys=["mp_id"])
         self._target_keys_field = "xas_ids"
 
-    def process_item(self, spectra: List[Dict]) -> Dict:
+    def process_item(self, spectra: list[dict]) -> dict:
         # TODO: Change this to do structure matching against materials collection
         mpid = spectra[0]["mp_id"]
 

emmet/builders/materials/absorption_spectrum.py

@@ -1,5 +1,5 @@
+from __future__ import annotations
 from math import ceil
-from typing import Dict, Iterator, List, Optional
 
 import numpy as np
 from maggma.builders import Builder
@@ -10,6 +10,11 @@ from pymatgen.core.structure import Structure
 from emmet.core.absorption import AbsorptionDoc
 from emmet.core.utils import jsanitize
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 
 class AbsorptionBuilder(Builder):
     def __init__(
@@ -17,7 +22,7 @@ class AbsorptionBuilder(Builder):
         materials: Store,
         tasks: Store,
         absorption: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         self.materials = materials
@@ -32,7 +37,7 @@ class AbsorptionBuilder(Builder):
 
         super().__init__(sources=[materials, tasks], targets=[absorption], **kwargs)
 
-    def prechunk(self, number_splits: int) -> Iterator[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterator[dict]:  # pragma: no cover
         """
         Prechunk method to perform chunking by the key field
         """
@@ -44,7 +49,7 @@ class AbsorptionBuilder(Builder):
         for split in grouper(keys, N):
             yield {"query": {self.materials.key: {"$in": list(split)}}}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets all items to process
 

emmet/builders/materials/alloys.py

@@ -1,5 +1,4 @@
 from itertools import chain, combinations
-from typing import Dict, List, Tuple, Union
 
 from maggma.builders import Builder
 from matminer.datasets import load_dataset
@@ -40,7 +39,7 @@ class AlloyPairBuilder(Builder):
         provenance,
         oxi_states,
         alloy_pairs,
-        thermo_type: Union[ThermoType, str] = ThermoType.GGA_GGA_U_R2SCAN,
+        thermo_type: ThermoType | str = ThermoType.GGA_GGA_U_R2SCAN,
     ):
         self.materials = materials
         self.thermo = thermo
@@ -274,7 +273,7 @@ class AlloyPairMemberBuilder(Builder):
             if structures:
                 yield (pairs, structures)
 
-    def process_item(self, item: Tuple[List[AlloyPair], Dict[str, Structure]]):
+    def process_item(self, item: tuple[list[AlloyPair], dict[str, Structure]]):
         pairs, structures = item
 
         all_pair_members = []

emmet/builders/materials/chemenv.py

@@ -1,8 +1,7 @@
-from typing import Dict, Optional
 from maggma.builders.map_builder import MapBuilder
 from maggma.core import Store
-
 from pymatgen.core.structure import Structure
+
 from emmet.core.chemenv import ChemEnvDoc
 from emmet.core.utils import jsanitize
 
@@ -12,7 +11,7 @@ class ChemEnvBuilder(MapBuilder):
         self,
         oxidation_states: Store,
         chemenv: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs
     ):
         self.oxidation_states = oxidation_states

emmet/builders/materials/corrected_entries.py

@@ -1,10 +1,11 @@
+from __future__ import annotations
+
 import copy
 import warnings
 from collections import defaultdict
 from datetime import datetime
 from itertools import chain
 from math import ceil
-from typing import Dict, Iterable, Iterator, List, Optional, Union
 
 from maggma.core import Builder, Store
 from maggma.utils import grouper
@@ -16,15 +17,20 @@ from emmet.core.corrected_entries import CorrectedEntriesDoc
 from emmet.core.thermo import ThermoType
 from emmet.core.utils import jsanitize
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator
+
 
 class CorrectedEntriesBuilder(Builder):
     def __init__(
         self,
         materials: Store,
         corrected_entries: Store,
-        oxidation_states: Optional[Store] = None,
-        query: Optional[Dict] = None,
-        compatibility: Optional[Union[List[Compatibility], List[None]]] = [None],
+        oxidation_states: Store | None = None,
+        query: dict | None = None,
+        compatibility: list[Compatibility] | list[None] | None = [None],
         chunk_size: int = 1000,
         **kwargs,
     ):
@@ -48,7 +54,7 @@ class CorrectedEntriesBuilder(Builder):
         self.compatibility = compatibility
         self.oxidation_states = oxidation_states
         self.chunk_size = chunk_size
-        self._entries_cache: Dict[str, List[Dict]] = defaultdict(list)
+        self._entries_cache: dict[str, list[dict]] = defaultdict(list)
 
         if self.corrected_entries.key != "chemsys":
             warnings.warn(
@@ -93,7 +99,7 @@ class CorrectedEntriesBuilder(Builder):
         # Search index for corrected_entries
         self.corrected_entries.ensure_index("chemsys")
 
-    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
         to_process_chemsys = self._get_chemsys_to_process()
 
         N = ceil(len(to_process_chemsys) / number_splits)
@@ -101,7 +107,7 @@ class CorrectedEntriesBuilder(Builder):
         for chemsys_chunk in grouper(to_process_chemsys, N):
             yield {"query": {"chemsys": {"$in": list(chemsys_chunk)}}}
 
-    def get_items(self) -> Iterator[List[Dict]]:
+    def get_items(self) -> Iterator[list[dict]]:
         """
         Gets whole chemical systems of entries to process
         """
@@ -210,7 +216,7 @@ class CorrectedEntriesBuilder(Builder):
         else:
             self.logger.info("No corrected entry items to update")
 
-    def get_entries(self, chemsys: str) -> List[Dict]:
+    def get_entries(self, chemsys: str) -> list[dict]:
         """
         Gets entries from the materials collection for the corresponding chemical systems
         Args:

emmet/builders/materials/dielectric.py

@@ -1,6 +1,6 @@
-from math import ceil
-from typing import Dict, Optional, Iterator
+from __future__ import annotations
 
+from math import ceil
 
 import numpy as np
 from maggma.builders import Builder
@@ -11,6 +11,11 @@ from pymatgen.core.structure import Structure
 from emmet.core.polar import DielectricDoc
 from emmet.core.utils import jsanitize
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 
 class DielectricBuilder(Builder):
     def __init__(
@@ -18,7 +23,7 @@ class DielectricBuilder(Builder):
         materials: Store,
         tasks: Store,
         dielectric: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         **kwargs,
     ):
         self.materials = materials
@@ -33,7 +38,7 @@ class DielectricBuilder(Builder):
 
         super().__init__(sources=[materials, tasks], targets=[dielectric], **kwargs)
 
-    def prechunk(self, number_splits: int) -> Iterator[Dict]:  # pragma: no cover
+    def prechunk(self, number_splits: int) -> Iterator[dict]:  # pragma: no cover
         """
         Prechunk method to perform chunking by the key field
         """

emmet/builders/materials/elasticity.py

@@ -17,8 +17,9 @@ The build proceeds in the below steps:
 7. Fit the elastic tensor.
 """
 
+from __future__ import annotations
+
 from datetime import datetime
-from typing import Any, Dict, Generator, List, Optional, Tuple, Union
 
 import numpy as np
 from maggma.core import Builder, Store
@@ -33,6 +34,12 @@ from emmet.core.mpid import MPID
 from emmet.core.utils import jsanitize
 from emmet.core.vasp.calc_types import CalcType
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+    from typing import Any
+
 
 class ElasticityBuilder(Builder):
     def __init__(
@@ -40,7 +47,7 @@ class ElasticityBuilder(Builder):
         tasks: Store,
         materials: Store,
         elasticity: Store,
-        query: Optional[Dict] = None,
+        query: dict | None = None,
         fitting_method: str = "finite_difference",
         **kwargs,
     ):
@@ -78,7 +85,7 @@ class ElasticityBuilder(Builder):
 
     def get_items(
         self,
-    ) -> Generator[Tuple[str, Dict[str, str], List[Dict]], None, None]:
+    ) -> Generator[tuple[str, dict[str, str], list[dict]], None, None]:
         """
         Gets all items to process into elasticity docs.
 
@@ -129,8 +136,8 @@ class ElasticityBuilder(Builder):
             yield material_id, calc_types, tasks
 
     def process_item(
-        self, item: Tuple[MPID, Dict[str, str], List[Dict]]
-    ) -> Union[Dict, None]:
+        self, item: tuple[MPID, dict[str, str], list[dict]]
+    ) -> dict | None:
         """
         Process all tasks belong to the same material into an elasticity doc.
 
@@ -220,7 +227,7 @@ class ElasticityBuilder(Builder):
 
         return elasticity_doc
 
-    def update_targets(self, items: List[Dict]):
+    def update_targets(self, items: list[dict]):
         """
         Insert the new elasticity docs into the elasticity collection.
 
@@ -233,10 +240,10 @@ class ElasticityBuilder(Builder):
 
 
 def filter_opt_tasks(
-    tasks: List[Dict],
-    calc_types: Dict[str, str],
+    tasks: list[dict],
+    calc_types: dict[str, str],
     target_calc_type: str = CalcType.GGA_Structure_Optimization,
-) -> List[Dict]:
+) -> list[dict]:
     """
     Filter optimization tasks, by
         - calculation type
@@ -247,10 +254,10 @@ def filter_opt_tasks(
 
 
 def filter_deform_tasks(
-    tasks: List[Dict],
-    calc_types: Dict[str, str],
+    tasks: list[dict],
+    calc_types: dict[str, str],
     target_calc_type: str = CalcType.GGA_Deformation,
-) -> List[Dict]:
+) -> list[dict]:
     """
     Filter deformation tasks, by
         - calculation type
@@ -271,8 +278,8 @@ def filter_deform_tasks(
 
 
 def filter_by_incar_settings(
-    tasks: List[Dict], incar_settings: Optional[Dict[str, Any]] = None
-) -> List[Dict]:
+    tasks: list[dict], incar_settings: dict[str, Any] | None = None
+) -> list[dict]:
     """
     Filter tasks by incar parameters.
     """
@@ -315,7 +322,7 @@ def filter_by_incar_settings(
     return selected
 
 
-def filter_opt_tasks_by_time(tasks: List[Dict], logger) -> Dict:
+def filter_opt_tasks_by_time(tasks: list[dict], logger) -> dict:
     """
     Filter a set of tasks to select the latest completed one.
 
@@ -330,8 +337,8 @@ def filter_opt_tasks_by_time(tasks: List[Dict], logger) -> Dict:
 
 
 def filter_deform_tasks_by_time(
-    tasks: List[Dict], deform_comp_tol: float = 1e-5, logger=None
-) -> List[Dict]:
+    tasks: list[dict], deform_comp_tol: float = 1e-5, logger=None
+) -> list[dict]:
     """
     For deformation tasks with the same deformation, select the latest completed one.
 
@@ -364,7 +371,7 @@ def filter_deform_tasks_by_time(
     return selected
 
 
-def _filter_tasks_by_time(tasks: List[Dict], mode: str, logger) -> Dict:
+def _filter_tasks_by_time(tasks: list[dict], mode: str, logger) -> dict:
     """
     Helper function to filter a set of tasks to select the latest completed one.
     """
@@ -388,11 +395,11 @@ def _filter_tasks_by_time(tasks: List[Dict], mode: str, logger) -> Dict:
 
 
 def select_final_opt_deform_tasks(
-    opt_tasks: List[Tuple[np.ndarray, Dict]],
-    deform_tasks: List[Tuple[np.ndarray, List[Dict]]],
+    opt_tasks: list[tuple[np.ndarray, dict]],
+    deform_tasks: list[tuple[np.ndarray, list[dict]]],
     logger,
     lattice_comp_tol: float = 1e-5,
-) -> Tuple[Union[Dict, None], Union[List[Dict], None]]:
+) -> tuple[dict | None, list[dict] | None]:
     """
     Select the final opt task and deform tasks for fitting.
 
@@ -445,8 +452,8 @@ def select_final_opt_deform_tasks(
 
 
 def group_by_parent_lattice(
-    tasks: List[Dict], mode: str, lattice_comp_tol: float = 1e-5
-) -> List[Tuple[np.ndarray, List[Dict]]]:
+    tasks: list[dict], mode: str, lattice_comp_tol: float = 1e-5
+) -> list[tuple[np.ndarray, list[dict]]]:
     """
     Groups a set of task docs by parent lattice equivalence.
 
@@ -459,11 +466,11 @@ def group_by_parent_lattice(
         lattice_comp_tol: tolerance for comparing lattice equivalence.
 
     Returns:
-        [(lattice, List[tasks])]: each tuple gives the common parent lattice of a
+        [(lattice, list[tasks])]: each tuple gives the common parent lattice of a
             list of the structures before deformation (if any), and the list tasks
             from which the structures are taken.
     """
-    docs_by_lattice: List[Tuple[np.ndarray, List[Dict]]] = []
+    docs_by_lattice: list[tuple[np.ndarray, list[dict]]] = []
 
     for doc in tasks:
         sim_lattice = get(doc, "output.structure.lattice.matrix")