Update type checking.

shyuep · shyuep · commit 294134b67703 · 2025-04-27T16:42:59.000-07:00
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
@@ -31,7 +31,7 @@ jobs:
           uv run ruff format --check .
 
       - name: mypy
-        run: uv run mypy -p pymatgen.core
+        run: uv run mypy -p pymatgen.core -p pymatgen.symmetry -p pymatgen.util
 
       - name: pyright
         run: uv run pyright src
diff --git a/pyproject.toml b/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"] #, "operator", "arg-type", "index", "call-arg", "return-value", "assignment", "attr-defined"]
-exclude = ['src/pymatgen/analysis', 'src/pymatgen/io', 'src/pymatgen/cli', 'src/pymatgen/electronic_structure', 'src/pymatgen/phonon']
+exclude = ['src/pymatgen/analysis', 'src/pymatgen/io', 'src/pymatgen/cli', 'src/pymatgen/electronic_structure', 'src/pymatgen/phonon', 'src/pymatgen/apps', 'src/pymatgen/vis', "src/pymatgen/alchemy"]
 plugins = ["numpy.typing.mypy_plugin"]
 
 [[tool.mypy.overrides]]
diff --git a/src/pymatgen/command_line/bader_caller.py b/src/pymatgen/command_line/bader_caller.py
@@ -28,6 +28,7 @@
 from monty.shutil import decompress_file
 from monty.tempfile import ScratchDir
 
+from pymatgen.core import Structure
 from pymatgen.io.common import VolumetricData
 from pymatgen.io.vasp.inputs import Potcar
 from pymatgen.io.vasp.outputs import Chgcar
@@ -37,7 +38,7 @@
 
     from typing_extensions import Self
 
-    from pymatgen.core import IStructure, Structure
+    from pymatgen.core import IStructure
 
 __author__ = "shyuepingong"
 __version__ = "0.1"
@@ -375,7 +376,7 @@ def get_oxidation_state_decorated_structure(self, nelects: list[int] | None = No
         Returns:
             Structure: with bader-analysis-based oxidation states.
         """
-        struct = self.structure.copy()
+        struct = Structure.from_sites(self.structure)
         charges = [self.get_partial_charge(idx, None if not nelects else nelects[idx]) for idx in range(len(struct))]
         struct.add_oxidation_state_by_site(charges)
         return struct
@@ -407,12 +408,12 @@ def get_decorated_structure(self, property_name: str, average: bool = False) ->
             structure with site properties assigned via Bader Analysis
         """
         vals = np.array([self.get_charge(i) for i in range(len(self.structure))])
-        struct = self.structure.copy()
+        struct = Structure.from_sites(self.structure)
         if average:
             vals = np.divide(vals, [d["atomic_vol"] for d in self.data])
         struct.add_site_property(property_name, vals)
         if property_name == "spin":
-            struct.add_spin_by_site(vals)
+            struct.add_spin_by_site(list(vals))
         return struct
 
     @property
diff --git a/src/pymatgen/command_line/chargemol_caller.py b/src/pymatgen/command_line/chargemol_caller.py
@@ -61,7 +61,7 @@
     from pathlib import Path
     from typing import Literal
 
-    from pymatgen.core import Structure
+    from pymatgen.core import IStructure, Structure
 
 __author__ = "Martin Siron, Andrew S. Rosen"
 __version__ = "0.1"
@@ -121,7 +121,7 @@ def __init__(
 
         if self._chgcar_path:
             self.chgcar: Chgcar | None = Chgcar.from_file(self._chgcar_path)
-            self.structure: Structure | None = self.chgcar.structure
+            self.structure: Structure | IStructure | None = self.chgcar.structure
             self.natoms: list[int] | None = self.chgcar.poscar.natoms
 
         else:
diff --git a/src/pymatgen/command_line/critic2_caller.py b/src/pymatgen/command_line/critic2_caller.py
@@ -81,7 +81,7 @@ class Critic2Caller:
     """Call critic2 and store standard output for further processing."""
 
     @requires(
-        which("critic2"),
+        which("critic2"),  # type: ignore[arg-type]
         "Critic2Caller requires the executable critic to be in the path. "
         "Please follow the instructions at https://github.com/aoterodelaroza/critic2.",
     )
diff --git a/src/pymatgen/command_line/enumlib_caller.py b/src/pymatgen/command_line/enumlib_caller.py
@@ -58,7 +58,7 @@
 
 
 @requires(
-    ENUM_CMD and MAKESTR_CMD,
+    ENUM_CMD and MAKESTR_CMD,  # type: ignore[arg-type]
     "EnumlibAdaptor requires the executables 'enum.x' or 'multienum.x' "
     "and 'makestr.x' or 'makeStr.py' to be in the path. Please download the "
     "library at https://github.com/msg-byu/enumlib and follow the instructions "
@@ -120,7 +120,7 @@ def __init__(
             finder = SpacegroupAnalyzer(structure, symm_prec)
             self.structure = finder.get_refined_structure()
         else:
-            self.structure = structure
+            self.structure = structure  # type: ignore[assignment]
 
         self.min_cell_size = min_cell_size
         self.max_cell_size = max_cell_size
@@ -372,7 +372,7 @@ def _get_structures(self, num_structs: int) -> list[Structure]:
                 data = _file.read()
                 data = re.sub(r"scale factor", "1", data)
                 data = re.sub(r"(\d+)-(\d+)", r"\1 -\2", data)
-                poscar = Poscar.from_str(data, self.index_species)
+                poscar = Poscar.from_str(data, self.index_species)  # type: ignore[arg-type]
                 sub_structure = poscar.structure
                 # Enumeration may have resulted in a super lattice. We need to
                 # find the mapping from the new lattice to the old lattice, and
diff --git a/src/pymatgen/command_line/mcsqs_caller.py b/src/pymatgen/command_line/mcsqs_caller.py
@@ -31,7 +31,7 @@ class Sqs(NamedTuple):
 
 
 @requires(
-    which("mcsqs") and which("str2cif"),
+    which("mcsqs") and which("str2cif"),  # type: ignore[arg-type]
     "run_mcsqs requires first installing AT-AT, see https://www.brown.edu/Departments/Engineering/Labs/avdw/atat/",
 )
 def run_mcsqs(
diff --git a/src/pymatgen/command_line/vampire_caller.py b/src/pymatgen/command_line/vampire_caller.py
@@ -53,7 +53,7 @@ class VampireCaller:
     """
 
     @requires(
-        VAMP_EXE,
+        VAMP_EXE,  # type: ignore[arg-type]
         "VampireCaller requires vampire-serial to be in the path."
         "Please follow the instructions at https://vampire.york.ac.uk/download/.",
     )
diff --git a/src/pymatgen/core/operations.py b/src/pymatgen/core/operations.py
@@ -117,7 +117,7 @@ def from_rotation_and_translation(
         affine_matrix[:3][:, 3] = translation_vec
         return cls(affine_matrix, tol)
 
-    def operate(self, point: NDArray[np.float64]) -> NDArray[np.float64]:
+    def operate(self, point: ArrayLike) -> NDArray[np.float64]:
         """Apply the operation on a point.
 
         Args:
@@ -129,7 +129,7 @@ def operate(self, point: NDArray[np.float64]) -> NDArray[np.float64]:
         affine_point = np.asarray([*point, 1])
         return np.dot(self.affine_matrix, affine_point)[:3]
 
-    def operate_multi(self, points: NDArray[np.float64]) -> NDArray[np.float64]:
+    def operate_multi(self, points: ArrayLike) -> NDArray[np.float64]:
         """Apply the operation on a list of points.
 
         Args:
diff --git a/src/pymatgen/core/structure.py b/src/pymatgen/core/structure.py
@@ -5063,7 +5063,7 @@ class Molecule(IMolecule, collections.abc.MutableSequence):
 
     def __init__(
         self,
-        species: Sequence[SpeciesLike],
+        species: Sequence[CompositionLike],
         coords: Sequence[ArrayLike] | ArrayLike,
         charge: float = 0.0,
         spin_multiplicity: int | None = None,
diff --git a/src/pymatgen/entries/compatibility.py b/src/pymatgen/entries/compatibility.py
@@ -1439,7 +1439,7 @@ def __init__(
         self.name = "MP Aqueous free energy adjustment"
         super().__init__()
 
-    def get_adjustments(self, entry: ComputedEntry) -> list[TemperatureEnergyAdjustment]:
+    def get_adjustments(self, entry: ComputedEntry) -> list[EnergyAdjustment]:
         """Get the corrections applied to a particular entry.
 
         Args:
@@ -1452,7 +1452,7 @@ def get_adjustments(self, entry: ComputedEntry) -> list[TemperatureEnergyAdjustm
             CompatibilityError if the required O2 and H2O energies have not been provided to
             MaterialsProjectAqueousCompatibility during init or in the list of entries passed to process_entries.
         """
-        adjustments = []
+        adjustments: list[EnergyAdjustment] = []
         if self.o2_energy is None or self.h2o_energy is None or self.h2o_adjustments is None:
             raise CompatibilityError(
                 "You did not provide the required O2 and H2O energies. "
diff --git a/src/pymatgen/entries/entry_tools.py b/src/pymatgen/entries/entry_tools.py
@@ -29,6 +29,7 @@
 
     from pymatgen.entries import Entry
     from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
+    from pymatgen.util.typing import SpeciesLike
 
 logger = logging.getLogger(__name__)
 
@@ -309,7 +310,7 @@ def to_csv(self, filename: str, latexify_names: bool = False) -> None:
             latexify_names: Format entry names to be LaTex compatible,
                 e.g. Li_{2}O
         """
-        els: set[Element] = set()
+        els: set[SpeciesLike] = set()
         for entry in self.entries:
             els.update(entry.elements)
         elements = sorted(els, key=lambda a: a.X)
diff --git a/src/pymatgen/entries/mixing_scheme.py b/src/pymatgen/entries/mixing_scheme.py
@@ -157,7 +157,7 @@ def process_entries(
         processed_entry_list: list = []
 
         # We can't operate on single entries in this scheme
-        if len(entries) == 1:
+        if len(entries) == 1:  # type: ignore[arg-type]
             warnings.warn(
                 f"{type(self).__name__} cannot process single entries. Supply a list of entries.", stacklevel=2
             )
@@ -500,12 +500,12 @@ def get_mixing_state_data(self, entries: list[ComputedStructureEntry]):
         # construct PhaseDiagram for each run_type, if possible
         pd_type_1, pd_type_2 = None, None
         try:
-            pd_type_1 = PhaseDiagram(entries_type_1)
+            pd_type_1 = PhaseDiagram(entries_type_1)  # type: ignore[arg-type]
         except ValueError:
             warnings.warn(f"{self.run_type_1} entries do not form a complete PhaseDiagram.", stacklevel=2)
 
         try:
-            pd_type_2 = PhaseDiagram(entries_type_2)
+            pd_type_2 = PhaseDiagram(entries_type_2)  # type: ignore[arg-type]
         except ValueError:
             warnings.warn(f"{self.run_type_2} entries do not form a complete PhaseDiagram.", stacklevel=2)
 
diff --git a/src/pymatgen/ext/optimade.py b/src/pymatgen/ext/optimade.py
@@ -19,6 +19,8 @@
 
     from typing_extensions import Self
 
+    from pymatgen.core.structure import Molecule
+
 _logger = logging.getLogger(__name__)
 _handler = logging.StreamHandler(sys.stdout)
 _logger.addHandler(_handler)
@@ -235,7 +237,7 @@ def get_structures(
         nsites: int | None = None,
         chemical_formula_anonymous: str | None = None,
         chemical_formula_hill: str | None = None,
-    ) -> dict[str, dict[str, Structure]]:
+    ) -> dict[str, dict[str, Structure | Molecule]]:
         """Retrieve Structures from OPTIMADE providers.
 
         Not all functionality of OPTIMADE is currently exposed in this convenience method. To
@@ -307,7 +309,7 @@ def get_snls(
 
         return self.get_snls_with_filter(optimade_filter, additional_response_fields=additional_response_fields)
 
-    def get_structures_with_filter(self, optimade_filter: str) -> dict[str, dict[str, Structure]]:
+    def get_structures_with_filter(self, optimade_filter: str) -> dict[str, dict[str, Structure | Molecule]]:
         """Get structures satisfying a given OPTIMADE filter.
 
         Args:
@@ -416,7 +418,7 @@ def _get_comp(sp_dict):
                 # TODO: follow `references` to add reference information here
                 snl = StructureNL(
                     structure,
-                    authors={},
+                    authors=[],
                     history=[
                         {
                             "name": identifier,
@@ -455,7 +457,7 @@ def _get_comp(sp_dict):
                     # TODO: follow `references` to add reference information here
                     snl = StructureNL(
                         structure,
-                        authors={},
+                        authors=[],
                         history=[
                             {
                                 "name": identifier,
diff --git a/src/pymatgen/symmetry/analyzer.py b/src/pymatgen/symmetry/analyzer.py
@@ -90,7 +90,7 @@ class SpacegroupAnalyzer:
     def __init__(
         self,
         structure: Structure | IStructure,
-        symprec: float | None = 0.01,
+        symprec: float = 0.01,
         angle_tolerance: float = 5,
     ) -> None:
         """
@@ -197,7 +197,7 @@ def get_point_group_symbol(self) -> str:
         # passing a 0-length rotations list to spglib can segfault
         if len(rotations) == 0:
             return "1"
-        return spglib.get_pointgroup(rotations)[0].strip()
+        return spglib.get_pointgroup(rotations)[0].strip()  # type:ignore[index]
 
     def get_crystal_system(self) -> CrystalSystem:
         """Get the crystal system for the structure, e.g. (triclinic, orthorhombic,
@@ -573,9 +573,9 @@ def get_primitive_standard_structure(
             international_monoclinic=international_monoclinic,
             keep_site_properties=keep_site_properties,
         )
-        lattice = self.get_lattice_type()
+        lattype = self.get_lattice_type()
 
-        if "P" in self.get_space_group_symbol() or lattice == "hexagonal":
+        if "P" in self.get_space_group_symbol() or lattype == "hexagonal":
             return conv
 
         transf = self.get_conventional_to_primitive_transformation_matrix(
@@ -861,15 +861,15 @@ def get_conventional_standard_structure(
                 # The following will convert to proper international convention
                 # that beta is the non-right angle.
                 op = [[0, 1, 0], [1, 0, 0], [0, 0, -1]]
-                transf = np.dot(op, transf)  # type: ignore[arg-type]
+                transf = np.dot(op, transf)
                 new_matrix = np.dot(op, new_matrix)
-                beta = Lattice(new_matrix).beta
+                beta = Lattice(new_matrix).beta  # type:ignore[arg-type]
                 if beta < 90:
                     op = [[-1, 0, 0], [0, -1, 0], [0, 0, 1]]
                     transf = np.dot(op, transf)  # type: ignore[arg-type]
                     new_matrix = np.dot(op, new_matrix)  # type: ignore[arg-type]
 
-            lattice = Lattice(new_matrix)
+            lattice = Lattice(new_matrix)  # type:ignore[arg-type]
 
         elif latt_type == "triclinic":
             # we use a LLL Minkowski-like reduction for the triclinic cells
@@ -954,7 +954,7 @@ def is_all_acute_or_obtuse(matrix) -> bool:
                 transf = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])  # type:ignore[assignment]
                 new_matrix = test_matrix
 
-            lattice = Lattice(new_matrix)
+            lattice = Lattice(new_matrix)  # type:ignore[arg-type]
 
         new_coords = np.dot(transf, np.transpose(struct.frac_coords)).T  # type: ignore[arg-type]
         new_struct = Structure(
@@ -998,7 +998,7 @@ def get_kpoint_weights(self, kpoints: Sequence[Kpoint], atol: float = 1e-5) -> l
                 shift.append(1)
 
         mapping, grid = spglib.get_ir_reciprocal_mesh(np.array(mesh), self._cell, is_shift=shift, symprec=self._symprec)
-        mapping = list(mapping)
+        mapping = list(mapping)  # type:ignore[assignment]
         grid = (np.array(grid) + np.array(shift) * (0.5, 0.5, 0.5)) / mesh
         weights = []
         mapped: dict[tuple, int] = defaultdict(int)
@@ -1623,7 +1623,7 @@ def cluster_sites(
     mol: Molecule,
     tol: float,
     give_only_index: bool = False,
-) -> tuple[Site | None, dict]:
+) -> tuple[int | Site | None, dict]:
     """Cluster sites based on distance and species type.
 
     Args:
diff --git a/src/pymatgen/symmetry/groups.py b/src/pymatgen/symmetry/groups.py
@@ -23,7 +23,7 @@
 if TYPE_CHECKING:
     from typing import ClassVar, Literal, TypeAlias
 
-    from numpy.typing import ArrayLike
+    from numpy.typing import ArrayLike, NDArray
     from typing_extensions import Self
 
     from pymatgen.core.lattice import Lattice
@@ -177,7 +177,7 @@ def symmetry_ops(self) -> set[SymmOp]:
         """
         return self._symmetry_ops
 
-    def _generate_full_symmetry_ops(self) -> list[SymmOp]:
+    def _generate_full_symmetry_ops(self) -> list[NDArray]:
         symm_ops = list(self.generators)
         new_ops = self.generators
         while len(new_ops) > 0:
@@ -521,7 +521,7 @@ def get_orbit_and_generators(self, p: ArrayLike, tol: float = 1e-5) -> tuple[lis
 
         orbit: list[np.ndarray] = [np.array(p, dtype=float)]
         identity = SymmOp.from_rotation_and_translation(np.eye(3), np.zeros(3))
-        generators: list[np.ndarray] = [identity]
+        generators: list[SymmOp] = [identity]
         for o in self.symmetry_ops:
             pp = o.operate(p)
             pp = np.mod(np.round(pp, decimals=10), 1)
diff --git a/src/pymatgen/symmetry/kpath.py b/src/pymatgen/symmetry/kpath.py
diff --git a/src/pymatgen/symmetry/structure.py b/src/pymatgen/symmetry/structure.py
diff --git a/src/pymatgen/util/plotting.py b/src/pymatgen/util/plotting.py
diff --git a/src/pymatgen/util/string.py b/src/pymatgen/util/string.py

Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,7 @@ class Critic2Caller:`
`81`	`81`	`"""Call critic2 and store standard output for further processing."""`
`82`	`82`
`83`	`83`	`@requires(`
`84`		`- which("critic2"),`
	`84`	`+ which("critic2"), # type: ignore[arg-type]`
`85`	`85`	`"Critic2Caller requires the executable critic to be in the path. "`
`86`	`86`	`"Please follow the instructions at https://github.com/aoterodelaroza/critic2.",`
`87`	`87`	`)`
Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,7 @@ class Sqs(NamedTuple):`
`31`	`31`
`32`	`32`
`33`	`33`	`@requires(`
`34`		`- which("mcsqs") and which("str2cif"),`
	`34`	`+ which("mcsqs") and which("str2cif"), # type: ignore[arg-type]`
`35`	`35`	`"run_mcsqs requires first installing AT-AT, see https://www.brown.edu/Departments/Engineering/Labs/avdw/atat/",`
`36`	`36`	`)`
`37`	`37`	`def run_mcsqs(`
Original file line number	Diff line number	Diff line change
`@@ -53,7 +53,7 @@ class VampireCaller:`
`53`	`53`	`"""`
`54`	`54`
`55`	`55`	`@requires(`
`56`		`- VAMP_EXE,`
	`56`	`+ VAMP_EXE, # type: ignore[arg-type]`
`57`	`57`	`"VampireCaller requires vampire-serial to be in the path."`
`58`	`58`	`"Please follow the instructions at https://vampire.york.ac.uk/download/.",`
`59`	`59`	`)`