More typing fixes.

Author: shyuep

pyproject.toml

@@ -301,12 +301,12 @@ exclude_also = [
 ignore_missing_imports = true
 namespace_packages = true
 no_implicit_optional = false
-disable_error_code = ["annotation-unchecked", "override", "operator", "attr-defined", "union-attr", "arg-type"] #, "operator", "arg-type", "index", "call-arg", "return-value", "assignment", "attr-defined"]
+disable_error_code = ["annotation-unchecked", "override", "operator", "attr-defined", "union-attr"] #, "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']
 plugins = ["numpy.typing.mypy_plugin"]
 
 [[tool.mypy.overrides]]
-module = ["requests.*", "tabulate.*"]
+module = ["requests.*", "tabulate.*", "monty.*"]
 ignore_missing_imports = true
 
 [tool.codespell]

src/pymatgen/command_line/bader_caller.py

@@ -37,7 +37,7 @@
 
     from typing_extensions import Self
 
-    from pymatgen.core import Structure
+    from pymatgen.core import IStructure, Structure
 
 __author__ = "shyuepingong"
 __version__ = "0.1"
@@ -351,7 +351,7 @@ def get_partial_charge(self, atom_index: int, nelect: int | None = None) -> floa
         """
         return -self.get_charge_transfer(atom_index, nelect)
 
-    def get_charge_decorated_structure(self) -> Structure:
+    def get_charge_decorated_structure(self) -> Structure | IStructure:
         """Get a charge decorated structure.
 
         Note, this assumes that the Bader analysis was correctly performed on a file

src/pymatgen/core/ion.py

@@ -173,7 +173,8 @@ def get_reduced_formula_and_factor(
                 comp = self.composition - nH2O * Composition("H2O")
 
         el_amt_dict = {k: round(v) for k, v in comp.get_el_amt_dict().items()}
-        formula, factor = reduce_formula(el_amt_dict, iupac_ordering=iupac_ordering)
+        formula, factor_int = reduce_formula(el_amt_dict, iupac_ordering=iupac_ordering)
+        factor = float(factor_int)
 
         # This line checks specifically that the contains an equal amount of O and H. When that is the case,
         # they should be displayed as "OH" rather than "HO".

src/pymatgen/core/periodic_table.py

@@ -1056,7 +1056,6 @@ def __init__(
             ValueError: If oxidation state passed both in symbol string and via
                 oxidation_state kwarg.
         """
-        symbol = str(symbol)  # First ensure that any symbol is converted to a string.
         if oxidation_state is not None and symbol.endswith(("+", "-")):
             raise ValueError(
                 f"Oxidation state should be specified either in {symbol=} or as {oxidation_state=}, not both."

src/pymatgen/core/structure.py

@@ -2961,7 +2961,7 @@ def to(self, filename: PathLike = "", fmt: FileFormats = "", **kwargs) -> str:
             json_str = json.dumps(self.as_dict(), **kwargs)
             if filename:
                 with zopen(filename, mode="wt", encoding="utf-8") as file:
-                    file.write(json_str)
+                    file.write(json_str)  # type:ignore[arg-type]
             return json_str
 
         elif fmt == "xsf" or fnmatch(filename.lower(), "*.xsf*"):
@@ -2970,7 +2970,7 @@ def to(self, filename: PathLike = "", fmt: FileFormats = "", **kwargs) -> str:
             res_str = XSF(self).to_str()
             if filename:
                 with zopen(filename, mode="wt", encoding="utf-8") as file:
-                    file.write(res_str)
+                    file.write(res_str)  # type:ignore[arg-type]
             return res_str
 
         elif (
@@ -2984,7 +2984,7 @@ def to(self, filename: PathLike = "", fmt: FileFormats = "", **kwargs) -> str:
             res_str = Mcsqs(self).to_str()
             if filename:
                 with zopen(filename, mode="wt", encoding="ascii") as file:
-                    file.write(res_str)
+                    file.write(res_str)  # type:ignore[arg-type]
             return res_str
 
         elif fmt == "prismatic" or fnmatch(filename, "*prismatic*"):
@@ -2999,7 +2999,7 @@ def to(self, filename: PathLike = "", fmt: FileFormats = "", **kwargs) -> str:
             yaml_str = str_io.getvalue()
             if filename:
                 with zopen(filename, mode="wt", encoding="utf-8") as file:
-                    file.write(yaml_str)
+                    file.write(yaml_str)  # type:ignore[arg-type]
             return yaml_str
 
         elif fmt == "aims" or fnmatch(filename, "geometry.in"):
@@ -3008,9 +3008,9 @@ def to(self, filename: PathLike = "", fmt: FileFormats = "", **kwargs) -> str:
             geom_in = AimsGeometryIn.from_structure(self)
             if filename:
                 with zopen(filename, mode="wt", encoding="utf-8") as file:
-                    file.write(geom_in.get_header(filename))
-                    file.write(geom_in.content)
-                    file.write("\n")
+                    file.write(geom_in.get_header(filename))  # type:ignore[arg-type]
+                    file.write(geom_in.content)  # type:ignore[arg-type]
+                    file.write("\n")  # type:ignore[arg-type]
             return geom_in.content
 
         # fleur support implemented in external namespace pkg https://github.com/JuDFTteam/pymatgen-io-fleur
@@ -3025,7 +3025,7 @@ def to(self, filename: PathLike = "", fmt: FileFormats = "", **kwargs) -> str:
             res_str = ResIO.structure_to_str(self)
             if filename:
                 with zopen(filename, mode="wt", encoding="utf-8") as file:
-                    file.write(res_str)
+                    file.write(res_str)  # type:ignore[arg-type]
             return res_str
 
         elif fmt == "pwmat" or fnmatch(filename.lower(), "*.pwmat") or fnmatch(filename.lower(), "*.config"):
@@ -4013,7 +4013,7 @@ def to(self, filename: str = "", fmt: str = "") -> str | None:
             json_str = json.dumps(self.as_dict())
             if filename:
                 with zopen(filename, mode="wt", encoding="utf-8") as file:
-                    file.write(json_str.encode("utf-8"))
+                    file.write(json_str)  # type:ignore[arg-type]
             return json_str
         elif fmt in {"yaml", "yml"} or fnmatch(filename, "*.yaml*") or fnmatch(filename, "*.yml*"):
             yaml = YAML()
@@ -4022,7 +4022,7 @@ def to(self, filename: str = "", fmt: str = "") -> str | None:
             yaml_str = str_io.getvalue()
             if filename:
                 with zopen(filename, mode="wt", encoding="utf-8") as file:
-                    file.write(yaml_str.encode("utf-8"))
+                    file.write(yaml_str)  # type:ignore[arg-type]
             return yaml_str
         else:
             from pymatgen.io.babel import BabelMolAdaptor

src/pymatgen/core/trajectory.py

@@ -475,7 +475,7 @@ def write_Xdatcar(
         xdatcar_str = "\n".join(lines) + "\n"
 
         with zopen(filename, mode="wt", encoding="utf-8") as file:
-            file.write(xdatcar_str)
+            file.write(xdatcar_str)  # type:ignore[arg-type]
 
     def as_dict(self) -> dict:
         """Return the trajectory as a MSONable dict."""

src/pymatgen/entries/compatibility.py

@@ -773,7 +773,7 @@ def __init__(
         self.run_types = run_types
         super().__init__()
 
-    def get_adjustments(self, entry: AnyComputedEntry) -> list[EnergyAdjustment]:
+    def get_adjustments(self, entry: AnyComputedEntry) -> list[ConstantEnergyAdjustment]:
         """Get the list of energy adjustments to be applied to an entry."""
         adjustment_list = []
         if entry.parameters.get("run_type") not in self.run_types:
@@ -1053,7 +1053,7 @@ def __init__(
             self.u_corrections = {}
             self.u_errors = {}
 
-    def get_adjustments(self, entry: AnyComputedEntry) -> list[EnergyAdjustment]:
+    def get_adjustments(self, entry: AnyComputedEntry) -> list[CompositionEnergyAdjustment]:
         """Get the energy adjustments for a ComputedEntry or ComputedStructureEntry.
 
         Energy corrections are implemented directly in this method instead of in
@@ -1439,7 +1439,7 @@ def __init__(
         self.name = "MP Aqueous free energy adjustment"
         super().__init__()
 
-    def get_adjustments(self, entry: ComputedEntry) -> list[EnergyAdjustment]:
+    def get_adjustments(self, entry: ComputedEntry) -> list[TemperatureEnergyAdjustment]:
         """Get the corrections applied to a particular entry.
 
         Args:

src/pymatgen/io/abinit/abiobjects.py

@@ -20,6 +20,8 @@
 
     from typing_extensions import Self
 
+    from pymatgen.core.structure import IStructure
+
 
 def lattice_from_abivars(cls=None, *args, **kwargs):
     """Get a `Lattice` object from a dictionary with the Abinit variables `acell`
@@ -185,7 +187,7 @@ def species_by_znucl(structure: Structure) -> list[Species]:
 
 
 def structure_to_abivars(
-    structure: Structure,
+    structure: Structure | IStructure,
     enforce_znucl: list | None = None,
     enforce_typat: list | None = None,
     **kwargs,

src/pymatgen/io/aims/sets/base.py

@@ -20,6 +20,7 @@
 if TYPE_CHECKING:
     from collections.abc import Iterable, Sequence
 
+    from pymatgen.core.structure import IStructure
     from pymatgen.util.typing import PathLike
 
 TMPDIR_NAME: str = "tmpdir"
@@ -374,7 +375,7 @@ def get_parameter_updates(
 
     def d2k(
         self,
-        structure: Structure,
+        structure: Structure | IStructure,
         kpt_density: float | tuple[float, float, float] = 5.0,
         even: bool = True,
     ) -> Iterable[float]:
@@ -395,7 +396,7 @@ def d2k(
         recip_cell = structure.lattice.inv_matrix.transpose()
         return self.d2k_recip_cell(recip_cell, structure.lattice.pbc, kpt_density, even)
 
-    def k2d(self, structure: Structure, k_grid: np.ndarray[int]):
+    def k2d(self, structure: Structure | IStructure, k_grid: np.ndarray[int]):
         """Generate the kpoint density in each direction from given k_grid.
 
         Args:

src/pymatgen/io/aims/sets/bs.py

@@ -13,14 +13,16 @@
     from collections.abc import Sequence
     from typing import Any
 
+    from pymatgen.core.structure import IStructure
+
 
 class _SegmentDict(TypedDict):
     coords: list[list[float]]
     labels: list[str]
     length: int
 
 
-def prepare_band_input(structure: Structure, density: float = 20):
+def prepare_band_input(structure: Structure | IStructure, density: float = 20):
     """Prepare the band information needed for the FHI-aims control.in file.
 
     Args: