More mypy fixes.

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/xr', 'src/pymatgen/io/pwmat', 'src/pymatgen/io/fiesta' , 'src/pymatgen/io/nwchem', 'src/pymatgen/io/phonopy', 'src/pymatgen/io/abinit', 'src/pymatgen/io/exciting']
+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']
 plugins = ["numpy.typing.mypy_plugin"]
 
 [[tool.mypy.overrides]]

src/pymatgen/io/exciting/inputs.py

@@ -72,13 +72,13 @@ def __init__(
             raise ValueError("Structure with partial occupancies cannot be converted into exciting input!")
 
     @property
-    def lockxyz(self) -> NDArray:
+    def lockxyz(self) -> NDArray | None:
         """Selective dynamics site properties."""
-        return self.structure.site_properties.get("selective_dynamics")
+        return self.structure.site_properties.get("selective_dynamics")  # type:ignore[return-value]
 
     @lockxyz.setter
-    def lockxyz(self, lockxyz: ArrayLike) -> NDArray:
-        self.structure.add_site_property("selective_dynamics", lockxyz)
+    def lockxyz(self, lockxyz: ArrayLike):
+        self.structure.add_site_property("selective_dynamics", np.array(lockxyz))
 
     @classmethod
     def from_str(cls, data: str) -> Self:
@@ -183,8 +183,8 @@ def from_file(cls, filename: PathLike) -> Self:
             ExcitingInput
         """
         with zopen(filename, mode="rt", encoding="utf-8") as file:
-            data = file.read().replace("\n", "")
-        return cls.from_str(data)
+            data = file.read().replace("\n", "")  # type:ignore[arg-type]
+        return cls.from_str(data)  # type:ignore[arg-type]
 
     def write_etree(
         self,

src/pymatgen/io/fiesta.py

@@ -533,7 +533,7 @@ def write_file(self, filename: str | Path) -> None:
             filename: Filename.
         """
         with zopen(filename, mode="wt", encoding="utf-8") as file:
-            file.write(str(self))
+            file.write(str(self))  # type:ignore[arg-type]
 
     def as_dict(self):
         """MSONable dict."""
@@ -565,7 +565,7 @@ def from_dict(cls, dct: dict) -> Self:
         )
 
     @classmethod
-    def from_str(cls, string_input: str) -> Self:
+    def from_str(cls, string_input: str) -> FiestaInput:
         """
         Read an FiestaInput from a string. Currently tested to work with
         files generated from this class itself.
@@ -700,7 +700,7 @@ def from_str(cls, string_input: str) -> Self:
         )
 
     @classmethod
-    def from_file(cls, filename: str | Path) -> Self:
+    def from_file(cls, filename: str | Path) -> FiestaInput:
         """
         Read an Fiesta input from a file. Currently tested to work with
         files generated from this class itself.
@@ -712,7 +712,7 @@ def from_file(cls, filename: str | Path) -> Self:
             FiestaInput object
         """
         with zopen(filename, mode="rt", encoding="utf-8") as file:
-            return cls.from_str(file.read())
+            return cls.from_str(file.read())  # type:ignore[arg-type]
 
 
 class FiestaOutput:

src/pymatgen/io/nwchem.py

@@ -535,7 +535,7 @@ def from_file(cls, filename: str | Path) -> Self:
             NwInput object
         """
         with zopen(filename, mode="rt", encoding="utf-8") as file:
-            return cls.from_str(file.read())
+            return cls.from_str(file.read())  # type:ignore[arg-type]
 
 
 class NwInputError(Exception):

src/pymatgen/io/phonopy.py

@@ -163,11 +163,11 @@ def get_ph_bs_symm_line_from_dict(bands_dict, has_nac=False, labels_dict=None) -
         if eig_q:
             eigen_displacements.append(eig_q)
 
-    q_pts = np.array(q_pts)
+    q_pts = np.array(q_pts)  # type:ignore[assignment]
     # transpose to match the convention in PhononBandStructure
-    frequencies = np.transpose(frequencies)
+    frequencies = np.transpose(frequencies)  # type:ignore[assignment]
     if eigen_displacements:
-        eigen_displacements = np.transpose(eigen_displacements, (1, 0, 2, 3))
+        eigen_displacements = np.transpose(eigen_displacements, (1, 0, 2, 3))  # type:ignore[assignment]
 
     rec_lattice = Lattice(bands_dict["reciprocal_lattice"])
 
@@ -338,7 +338,7 @@ def get_phonon_dos_from_fc(
     p_doses = dict(zip(structure, dos_raw[1], strict=True))
 
     total_dos = PhononDos(dos_raw[0], dos_raw[1].sum(axis=0))
-    return CompletePhononDos(structure, total_dos, p_doses)
+    return CompletePhononDos(structure, total_dos, p_doses)  # type:ignore[arg-type]
 
 
 @requires(Phonopy, "phonopy is required to calculate phonon band structures")
@@ -464,7 +464,7 @@ def get_gruneisenparameter(gruneisen_path, structure=None, structure_path=None)
     return GruneisenParameter(
         gruneisen=gruneisen_np,
         qpoints=q_pts_np,
-        multiplicities=multiplicities_np,
+        multiplicities=multiplicities_np,  # type:ignore[arg-type]
         frequencies=frequencies_np,
         structure=structure,
     )

src/pymatgen/io/pwmat/inputs.py

@@ -360,15 +360,15 @@ def get_virial(self) -> np.ndarray | None:
 class AtomConfig(MSONable):
     """Object for representing the data in a atom.config or final.config file."""
 
-    def __init__(self, structure: Structure | IStructure, sort_structure: bool = False):
+    def __init__(self, structure: IStructure | Structure, sort_structure: bool = False):
         """Initialization function.
 
         Args:
             structure (Structure): Structure object
             sort_structure (bool, optional): Whether to sort the structure. Useful if species
                 are not grouped properly together. Defaults to False.
         """
-        self.structure: Structure = structure
+        self.structure: Structure = Structure.from_sites(structure)
         if sort_structure:
             self.structure = self.structure.get_sorted_structure()
         elements_counter = dict(sorted(Counter(self.structure.species).items()))
@@ -396,7 +396,7 @@ def from_str(cls, data: str, mag: bool = False) -> Self:
         properties: dict[str, float] = {}
         structure = Structure(
             lattice=ac_extractor.get_lattice(),
-            species=ac_extractor.get_types(),
+            species=ac_extractor.get_types(),  # type:ignore[arg-type]
             coords=ac_extractor.get_coords().reshape(-1, 3),
             coords_are_cartesian=False,
             properties=properties,
@@ -420,7 +420,7 @@ def from_file(cls, filename: PathLike, mag: bool = False) -> Self:
             AtomConfig object.
         """
         with zopen(filename, mode="rt", encoding="utf-8") as file:
-            return cls.from_str(data=file.read(), mag=mag)
+            return cls.from_str(data=file.read(), mag=mag)  # type:ignore[arg-type]
 
     @classmethod
     def from_dict(cls, dct: dict) -> Self:
@@ -468,7 +468,7 @@ def get_str(self) -> str:
     def write_file(self, filename: PathLike, **kwargs):
         """Write AtomConfig to a file."""
         with zopen(filename, mode="wt", encoding="utf-8") as file:
-            file.write(self.get_str(**kwargs))
+            file.write(self.get_str(**kwargs))  # type:ignore[arg-type]
 
     def as_dict(self):
         """
@@ -696,4 +696,4 @@ def get_hsp_row_str(label: str, index: int, coordinate: float) -> str:
     def write_file(self, filename: PathLike):
         """Write HighSymmetryPoint to a file."""
         with zopen(filename, mode="wt", encoding="utf-8") as file:
-            file.write(self.get_str())
+            file.write(self.get_str())  # type:ignore[arg-type]

src/pymatgen/io/pwmat/outputs.py

@@ -258,7 +258,7 @@ def _parse_kpt(self) -> tuple[np.ndarray, np.ndarray, dict[str, np.ndarray]]:
         row_idx: int = LineLocator.locate_all_lines(self.filename, content)[0]
         kpts: np.ndarray = np.zeros((self._num_kpts, 3))
         kpts_weight: np.ndarray = np.zeros(self._num_kpts)
-        hsps: dict[str, np.array] = {}
+        hsps: dict[str, np.ndarray] = {}
         for ii in range(num_rows):
             #  0.00000     0.00000    0.00000     0.03704           G
             tmp_row_lst: list[str] = linecache.getline(str(self.filename), row_idx + ii + 1).split()

src/pymatgen/io/xr.py

@@ -69,7 +69,7 @@ def write_file(self, filename: str | Path) -> None:
             filename (str): name of the file to write to.
         """
         with zopen(filename, mode="wt", encoding="utf-8") as file:
-            file.write(str(self) + "\n")
+            file.write(str(self) + "\n")  # type:ignore[arg-type]
 
     @classmethod
     def from_str(cls, string: str, use_cores: bool = True, thresh: float = 1e-4) -> Self:
@@ -156,4 +156,4 @@ def from_file(cls, filename: str | Path, use_cores: bool = True, thresh: float =
                     file.
         """
         with zopen(filename, mode="rt", encoding="utf-8") as file:
-            return cls.from_str(file.read(), use_cores=use_cores, thresh=thresh)
+            return cls.from_str(file.read(), use_cores=use_cores, thresh=thresh)  # type:ignore[arg-type]