More mypy cleanup.

Author: shyuep

pyproject.toml

@@ -302,7 +302,7 @@ ignore_missing_imports = true
 namespace_packages = true
 no_implicit_optional = false
 disable_error_code = ["annotation-unchecked", "override", "operator", "attr-defined", "union-attr", "misc", "call-overload", "call-arg", "var-annotated"] #, "operator", "arg-type", "index", "call-arg", "return-value", "assignment", "attr-defined"]
-exclude = ['src/pymatgen/analysis', 'src/pymatgen/phonon', 'src/pymatgen/io/lobster', 'src/pymatgen/io/cp2k', 'src/pymatgen/io/aims', 'src/pymatgen/io/lammps', 'src/pymatgen/io/feff', 'src/pymatgen/io/abinit']
+exclude = ['src/pymatgen/analysis', 'src/pymatgen/phonon', 'src/pymatgen/io/lobster', 'src/pymatgen/io/cp2k', 'src/pymatgen/io/aims', 'src/pymatgen/io/lammps']
 plugins = ["numpy.typing.mypy_plugin"]
 
 [[tool.mypy.overrides]]

src/pymatgen/core/units.py

@@ -24,7 +24,7 @@
     from collections.abc import Iterator
     from typing import Any
 
-    from numpy.typing import NDArray
+    from numpy.typing import ArrayLike
     from typing_extensions import Self
 
 __author__ = "Shyue Ping Ong, Matteo Giantomassi"
@@ -511,7 +511,7 @@ class ArrayWithUnit(np.ndarray):
 
     def __new__(
         cls,
-        input_array: NDArray,
+        input_array: ArrayLike,
         unit: str | Unit,
         unit_type: str | None = None,
     ) -> Self:

src/pymatgen/io/abinit/abiobjects.py

@@ -20,6 +20,7 @@
 
     from typing_extensions import Self
 
+    from pymatgen.core.periodic_table import DummySpecies, Element
     from pymatgen.core.structure import IStructure
 
 
@@ -164,7 +165,7 @@ def structure_from_abivars(cls=None, *args, **kwargs) -> Structure:
     )
 
 
-def species_by_znucl(structure: Structure) -> list[Species]:
+def species_by_znucl(structure: IStructure | Structure) -> list[Species | Element | DummySpecies]:
     """Get list of unique specie found in structure **ordered according to sites**.
 
     Example:
@@ -235,7 +236,7 @@ def structure_to_abivars(
         types_of_specie = species_by_znucl(structure)
 
         znucl_type = [specie.number for specie in types_of_specie]
-        typat = np.zeros(n_atoms, int)
+        typat = np.zeros(n_atoms, int)  # type:ignore[assignment]
         for atm_idx, site in enumerate(structure):
             typat[atm_idx] = types_of_specie.index(site.specie) + 1
 
@@ -245,8 +246,8 @@ def structure_to_abivars(
 
     # Set small values to zero. This usually happens when the CIF file
     # does not give structure parameters with enough digits.
-    r_prim = np.where(np.abs(r_prim) > 1e-8, r_prim, 0.0)
-    x_red = np.where(np.abs(x_red) > 1e-8, x_red, 0.0)
+    r_prim = np.where(np.abs(r_prim) > 1e-8, r_prim, 0.0)  # type:ignore[assignment]
+    x_red = np.where(np.abs(x_red) > 1e-8, x_red, 0.0)  # type:ignore[assignment]
 
     # Info on atoms.
     dct = {

src/pymatgen/io/abinit/inputs.py

@@ -753,7 +753,7 @@ def __init__(
             pseudo_dir = os.path.abspath(pseudo_dir)
             if not os.path.isdir(pseudo_dir):
                 raise self.Error(f"Directory {pseudo_dir} does not exist")
-            pseudos = [os.path.join(pseudo_dir, p) for p in pseudos]
+            pseudos = [os.path.join(pseudo_dir, p) for p in pseudos]  # type:ignore[arg-type]
 
         try:
             self._pseudos = PseudoTable.as_table(pseudos).get_pseudos_for_structure(self.structure)

src/pymatgen/io/abinit/pseudos.py

@@ -92,13 +92,13 @@ class Pseudo(MSONable, abc.ABC):
     """
 
     @classmethod
-    def as_pseudo(cls, obj: Self | str) -> Self:
+    def as_pseudo(cls, obj: Pseudo | str) -> Pseudo:
         """Convert obj into a Pseudo.
 
         Args:
             obj (str | Pseudo): Path to the pseudo file or a Pseudo object.
         """
-        return obj if isinstance(obj, cls) else cls.from_file(obj)
+        return obj if isinstance(obj, cls) else cls.from_file(obj)  # type:ignore[arg-type]
 
     @classmethod
     def from_file(cls, filename: str) -> Self:
@@ -1558,9 +1558,9 @@ def from_dir(cls, top, exts=None, exclude_dirs="_*") -> Self | None:
             if exts is None:
                 exts = ("psp8",)
 
-            for pseudo in find_exts(top, exts, exclude_dirs=exclude_dirs):
+            for pseudo in find_exts(top, exts, exclude_dirs=exclude_dirs):  # type:ignore[assignment]
                 try:
-                    pseudos.append(Pseudo.from_file(pseudo))
+                    pseudos.append(Pseudo.from_file(pseudo))  # type:ignore[arg-type]
                 except Exception as exc:
                     logger.critical(f"Error in {pseudo}:\n{exc}")
 

src/pymatgen/io/feff/inputs.py

@@ -188,14 +188,14 @@ def __init__(
         # if Structure, check symmetry
         if isinstance(self.struct, Structure):
             self.periodic = True
-            sym = SpacegroupAnalyzer(struct, **self.spacegroup_analyzer_settings)
+            sym = SpacegroupAnalyzer(struct, **self.spacegroup_analyzer_settings)  # type:ignore[arg-type]
             data = sym.get_symmetry_dataset()
             self.space_number = data.number
             self.space_group = data.international
         # for Molecule, skip the symmetry check
         elif isinstance(self.struct, Molecule):
             self.periodic = False
-            self.space_number = self.space_group = None
+            self.space_number = self.space_group = None  # type:ignore[assignment]
         else:
             raise TypeError("'struct' argument must be a Structure or Molecule!")
         self.comment = comment or "None given"
@@ -247,7 +247,7 @@ def header_string_from_file(filename: str = "feff.inp"):
             Reads header string.
         """
         with zopen(filename, mode="rt", encoding="utf-8") as file:
-            lines = file.readlines()
+            lines: list[str] = file.readlines()  # type:ignore[assignment]
             feff_header_str = []
             ln = 0
 
@@ -669,7 +669,7 @@ def from_file(cls, filename: str = "feff.inp") -> Self:
             Tags
         """
         with zopen(filename, mode="rt", encoding="utf-8") as file:
-            lines = list(clean_lines(file.readlines()))
+            lines: list[str] = list(clean_lines(file.readlines()))  # type:ignore[arg-type]
         params = {}
         eels_params = []
         ieels = -1

src/pymatgen/io/feff/outputs.py

@@ -71,7 +71,8 @@ def from_file(cls, feff_inp_file: str = "feff.inp", ldos_file: str = "ldos") ->
             begin = 0
 
             with zopen(pot_inp, mode="rt", encoding="utf-8") as potfile:
-                for line in potfile:
+                line: str
+                for line in potfile:  # type:ignore[assignment]
                     if len(pot_read_end.findall(line)) > 0:
                         break
 
@@ -144,7 +145,7 @@ def from_file(cls, feff_inp_file: str = "feff.inp", ldos_file: str = "ldos") ->
         _t_dos: dict = {Spin.up: t_dos}
 
         dos = Dos(e_fermi, dos_energies, _t_dos)
-        complete_dos = CompleteDos(structure, dos, pdoss)
+        complete_dos = CompleteDos(structure, dos, pdoss)  # type:ignore[arg-type]
         charge_transfer = LDos.charge_transfer_from_file(feff_inp_file, ldos_file)
         return cls(complete_dos, charge_transfer)