Mypy directives and type hinting for [Core, Builders, API]

emmet-api/emmet/api/routes/materials/materials/query_operators.py

@@ -1,18 +1,19 @@
 from itertools import permutations
-from typing import Literal, Optional
+from typing import Dict, Literal, Optional
 
-from emmet.core.symmetry import CrystalSystem
 from fastapi import Body, HTTPException, Query
 from maggma.api.query_operator import QueryOperator
 from maggma.api.utils import STORE_PARAMS
-from emmet.api.routes.materials.materials.utils import (
-    formula_to_criteria,
-    chemsys_to_criteria,
-)
 from pymatgen.analysis.structure_matcher import ElementComparator, StructureMatcher
 from pymatgen.core.composition import Composition, CompositionError
 from pymatgen.core.periodic_table import Element
 from pymatgen.core.structure import Structure
+
+from emmet.api.routes.materials.materials.utils import (
+    chemsys_to_criteria,
+    formula_to_criteria,
+)
+from emmet.core.symmetry import CrystalSystem
 from emmet.core.vasp.calc_types import RunType
 
 
@@ -281,9 +282,9 @@ class FindStructureQuery(QueryOperator):
 
     def query(
         self,
-        structure: Structure = Body(
+        structure: Dict = Body(
             ...,
-            description="Pymatgen structure object to query with",
+            description="Dictionary representaion of Pymatgen structure object to query with",
         ),
         ltol: float = Query(
             0.2,

emmet-api/tests/materials/materials/test_query_operators.py

@@ -1,6 +1,9 @@
 import os
 
-import pytest
+from monty.serialization import dumpfn, loadfn
+from monty.tempfile import ScratchDir
+from pymatgen.core.structure import Structure
+
 from emmet.api.core.settings import MAPISettings
 from emmet.api.routes.materials.materials.query_operators import (
     ChemsysQuery,
@@ -14,9 +17,6 @@
     SymmetryQuery,
 )
 from emmet.core.symmetry import CrystalSystem
-from monty.serialization import dumpfn, loadfn
-from monty.tempfile import ScratchDir
-from pymatgen.core.structure import Structure
 
 
 def test_formula_query():
@@ -158,7 +158,6 @@ def test_multi_material_id_query():
         }
 
 
-@pytest.mark.skip(reason="Pmg cif reader issue")
 def test_find_structure_query():
     op = FindStructureQuery()
 

emmet-builders/emmet/builders/abinit/phonon.py

@@ -1,41 +1,39 @@
-import tempfile
 import os
+import tempfile
 from math import ceil
-from emmet.builders.settings import EmmetBuildSettings
+from typing import Dict, Iterator, List, Optional, Tuple
+
 import numpy as np
-from typing import Optional, Dict, List, Iterator, Tuple
+from abipy.abio.inputs import AnaddbInput
+from abipy.core.abinit_units import eV_to_THz
+from abipy.dfpt.anaddbnc import AnaddbNcFile
+from abipy.dfpt.ddb import AnaddbError, DdbFile, DielectricTensorGenerator
+from abipy.dfpt.phonons import PhononBands
+from abipy.flowtk.tasks import AnaddbTask, TaskManager
+from maggma.builders import Builder
+from maggma.core import Store
 from maggma.utils import grouper
-
+from pymatgen.core.structure import Structure
+from pymatgen.io.abinit.abiobjects import KSampling
 from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
 from pymatgen.phonon.dos import CompletePhononDos
 from pymatgen.phonon.ir_spectra import IRDielectricTensor
-from pymatgen.core.structure import Structure
-from pymatgen.io.abinit.abiobjects import KSampling
-from pymatgen.symmetry.bandstructure import HighSymmKpath
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
-from abipy.dfpt.anaddbnc import AnaddbNcFile
-from abipy.abio.inputs import AnaddbInput
-from abipy.flowtk.tasks import AnaddbTask, TaskManager
-from abipy.dfpt.ddb import AnaddbError, DielectricTensorGenerator, DdbFile
-from abipy.dfpt.phonons import PhononBands
-from abipy.core.abinit_units import eV_to_THz
-from maggma.builders import Builder
-from maggma.core import Store
+from pymatgen.symmetry.bandstructure import HighSymmKpath
 
+from emmet.builders.settings import EmmetBuildSettings
 from emmet.core.phonon import (
-    PhononWarnings,
-    ThermodynamicProperties,
     AbinitPhonon,
-    VibrationalEnergy,
-)
-from emmet.core.phonon import (
-    PhononDos,
+    Ddb,
     PhononBandStructure,
+    PhononDos,
+    PhononWarnings,
     PhononWebsiteBS,
-    Ddb,
     ThermalDisplacement,
+    ThermodynamicProperties,
+    VibrationalEnergy,
 )
-from emmet.core.polar import DielectricDoc, BornEffectiveCharges, IRDielectric
+from emmet.core.polar import BornEffectiveCharges, DielectricDoc, IRDielectric
 from emmet.core.utils import jsanitize
 
 SETTINGS = EmmetBuildSettings()
@@ -679,7 +677,7 @@ def get_pmg_bs(
         )
 
         ph_bs_sl = PhononBandStructureSymmLine(
-            qpoints=qpts,
+            qpoints=qpts,  # type: ignore[arg-type]
             frequencies=ph_freqs,
             lattice=structure.reciprocal_lattice,
             has_nac=phbands.non_anal_ph is not None,

emmet-builders/emmet/builders/materials/alloys.py

@@ -1,19 +1,19 @@
-from itertools import combinations, chain
-from typing import Tuple, List, Dict, Union
+from itertools import chain, combinations
+from typing import Dict, List, Tuple, Union
 
-from tqdm import tqdm
 from maggma.builders import Builder
-from pymatgen.core.structure import Structure
 from matminer.datasets import load_dataset
-from emmet.core.thermo import ThermoType
-
 from pymatgen.analysis.alloys.core import (
-    AlloyPair,
-    InvalidAlloy,
     KNOWN_ANON_FORMULAS,
     AlloyMember,
+    AlloyPair,
     AlloySystem,
+    InvalidAlloy,
 )
+from pymatgen.core.structure import Structure
+from tqdm import tqdm
+
+from emmet.core.thermo import ThermoType
 
 # rough sort of ANON_FORMULAS by "complexity"
 ANON_FORMULAS = sorted(KNOWN_ANON_FORMULAS, key=lambda af: len(af))
@@ -291,7 +291,7 @@ def process_item(self, item: Tuple[List[AlloyPair], Dict[str, Structure]]):
                             is_ordered=structure.is_ordered,
                             x=pair.get_x(structure.composition),
                         )
-                        pair_members["members"].append(member.as_dict())
+                        pair_members["members"].append(member.as_dict())  # type: ignore[attr-defined]
                 except Exception as exc:
                     print(f"Exception for {db_id}: {exc}")
             if pair_members["members"]:

emmet-builders/emmet/builders/materials/corrected_entries.py

@@ -1,20 +1,20 @@
+import copy
 import warnings
 from collections import defaultdict
+from datetime import datetime
 from itertools import chain
-from typing import Dict, Iterable, Iterator, List, Optional
 from math import ceil
-import copy
-from datetime import datetime
+from typing import Dict, Iterable, Iterator, List, Optional, Union
 
 from maggma.core import Builder, Store
 from maggma.utils import grouper
-from pymatgen.entries.computed_entries import ComputedStructureEntry
 from pymatgen.entries.compatibility import Compatibility
+from pymatgen.entries.computed_entries import ComputedStructureEntry
 
-from emmet.core.utils import jsanitize
-from emmet.builders.utils import chemsys_permutations, HiddenPrints
-from emmet.core.thermo import ThermoType
+from emmet.builders.utils import HiddenPrints, chemsys_permutations
 from emmet.core.corrected_entries import CorrectedEntriesDoc
+from emmet.core.thermo import ThermoType
+from emmet.core.utils import jsanitize
 
 
 class CorrectedEntriesBuilder(Builder):
@@ -24,7 +24,7 @@ def __init__(
         corrected_entries: Store,
         oxidation_states: Optional[Store] = None,
         query: Optional[Dict] = None,
-        compatibility: Optional[List[Compatibility]] = None,
+        compatibility: Optional[Union[List[Compatibility], List[None]]] = [None],
         chunk_size: int = 1000,
         **kwargs,
     ):
@@ -45,10 +45,10 @@ def __init__(
         self.materials = materials
         self.query = query if query else {}
         self.corrected_entries = corrected_entries
-        self.compatibility = compatibility or [None]
+        self.compatibility = compatibility
         self.oxidation_states = oxidation_states
         self.chunk_size = chunk_size
-        self._entries_cache: Dict[str, List[ComputedStructureEntry]] = defaultdict(list)
+        self._entries_cache: Dict[str, List[Dict]] = defaultdict(list)
 
         if self.corrected_entries.key != "chemsys":
             warnings.warn(

emmet-builders/emmet/builders/materials/electrodes.py

@@ -1,22 +1,22 @@
 import operator
+from collections import defaultdict
 from datetime import datetime
 from functools import lru_cache
 from itertools import chain
 from math import ceil
-from typing import Any, Iterator, Dict, List, Optional
-from collections import defaultdict
+from typing import Any, Dict, Iterator, List, Optional
 
 from maggma.builders import Builder
 from maggma.stores import MongoStore
 from maggma.utils import grouper
-from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
+from pymatgen.analysis.phase_diagram import Composition, PhaseDiagram
 from pymatgen.entries.compatibility import MaterialsProject2020Compatibility
-from pymatgen.analysis.phase_diagram import PhaseDiagram, Composition
+from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
 
-from emmet.core.electrode import InsertionElectrodeDoc, ConversionElectrodeDoc
+from emmet.builders.settings import EmmetBuildSettings
+from emmet.core.electrode import ConversionElectrodeDoc, InsertionElectrodeDoc
 from emmet.core.structure_group import StructureGroupDoc, _get_id_lexi
 from emmet.core.utils import jsanitize
-from emmet.builders.settings import EmmetBuildSettings
 
 
 def s_hash(el):
@@ -564,9 +564,11 @@ def process_item(self, item) -> Dict:
                 # Get lowest material_id with matching composition
                 material_ids = [
                     (
-                        lambda x: x.data["material_id"]
-                        if x.composition.reduced_formula == v[1].reduced_formula
-                        else None
+                        lambda x: (
+                            x.data["material_id"]  # type: ignore[attr-defined]
+                            if x.composition.reduced_formula == v[1].reduced_formula
+                            else None
+                        )
                     )(e)
                     for e in pd.entries
                 ]
@@ -597,7 +599,7 @@ def process_item(self, item) -> Dict:
             relevant_entry_data = []
             for e in pd.entries:
                 if e.composition == Composition(c):
-                    relevant_entry_data.append((e.energy_per_atom, e.entry_id))
+                    relevant_entry_data.append((e.energy_per_atom, e.entry_id))  # type: ignore[attr-defined]
             relevant_entry_data.sort(key=lambda x: x[0])
             entry_id_mapping[c] = relevant_entry_data[0][1]
 

emmet-builders/emmet/builders/materials/provenance.py

@@ -1,7 +1,7 @@
 from collections import defaultdict
-from typing import Dict, Iterable, List, Optional, Tuple
-from math import ceil
 from datetime import datetime
+from math import ceil
+from typing import Dict, Iterable, List, Optional, Tuple
 
 from maggma.core import Builder, Store
 from maggma.utils import grouper
@@ -168,7 +168,7 @@ def get_items(self) -> Tuple[List[Dict], List[Dict]]:  # type: ignore
             for snl in snls:
                 struc = Structure.from_dict(snl)
                 snl_sg = get_sg(struc)
-                struc.snl = SNLDict(**snl)
+                struc.snl = SNLDict(**snl)  # type: ignore[attr-defined]
                 snl_groups[snl_sg].append(struc)
 
             mat_sg = get_sg(Structure.from_dict(mat["structure"]))

emmet-builders/emmet/builders/utils.py

@@ -1,17 +1,19 @@
 from __future__ import annotations
-from typing import Set, Union, Any, Literal, Optional
-import sys
+
+import json
 import os
-from pathlib import Path
+import sys
 from gzip import GzipFile
-import orjson
-import json
 from io import BytesIO
-from monty.serialization import MontyDecoder
-from botocore.exceptions import ClientError
 from itertools import chain, combinations
-from pymatgen.core import Structure
+from pathlib import Path
+from typing import Any, Literal, Optional, Set, Union
+
+import orjson
+from botocore.exceptions import ClientError
+from monty.serialization import MontyDecoder
 from pymatgen.analysis.diffusion.neb.full_path_mapper import MigrationGraph
+from pymatgen.core import Structure
 from pymatgen.io.vasp.inputs import PotcarSingle
 
 from emmet.builders.settings import EmmetBuildSettings
@@ -275,16 +277,16 @@ def get_potcar_stats(
                 # fallback method for validation - use header hash and symbol
                 # note that the potcar_spec assigns PotcarSingle.symbol to "titel"
                 # whereas the ***correct*** field is `header`
-                summary_stats["titel"] = potcar.header
-                summary_stats["hash"] = potcar.md5_header_hash
-                summary_stats = [summary_stats]
+                summary_stats["titel"] = potcar.header  # type: ignore[assignment]
+                summary_stats["hash"] = potcar.md5_header_hash  # type: ignore[assignment]
+                summary_stats = [summary_stats]  # type: ignore[assignment]
 
             elif method == "pymatgen":
-                summary_stats = []
+                summary_stats = []  # type: ignore[assignment]
                 for _, entries in PotcarSingle._potcar_summary_stats[
                     functional
                 ].items():
-                    summary_stats += [
+                    summary_stats += [  # type: ignore[operator]
                         {**entry, "titel": None, "hash": None}
                         for entry in entries
                         if entry["symbol"] == potcar_symbol

emmet-core/emmet/core/alloys.py

@@ -33,5 +33,5 @@ class AlloySystemDoc(EmmetBaseModel):
     @classmethod
     def from_pair(cls, alloy_system: AlloySystem):
         # Too large to duplicate alloy pairs here.
-        alloy_system.alloy_pairs = None
+        alloy_system.alloy_pairs = None  # type: ignore[assignment]
         return cls(alloy_system=alloy_system, alloy_id=alloy_system.alloy_id)

emmet-core/emmet/core/chemenv.py

@@ -442,7 +442,7 @@ def from_structure(
             inequivalent_indices_cations = [
                 indices[0]
                 for indices in symm_struct.equivalent_indices
-                if valences[indices[0]] > 0.0
+                if valences[indices[0]] > 0.0  # type: ignore[operator]
             ]
 
             se = lgf.compute_structure_environments(