mypy: more checks with pandas-stubs

Author: amorison

pyproject.toml

@@ -61,7 +61,6 @@ module = [
     "h5py.*",
     "matplotlib.*",
     "mpl_toolkits.*",
-    "pandas.*",
     "pytest.*",
     "scipy.*",
 ]

stagpy/_step.py

@@ -501,7 +501,7 @@ def _rprofs(self) -> DataFrame:
     def __getitem__(self, name: str) -> Rprof:
         step = self.step
         if name in self._rprofs.columns:
-            rprof = self._rprofs[name].values
+            rprof = self._rprofs[name].to_numpy()
             rad = self.centers
             if name in phyvars.RPROF:
                 meta = phyvars.RPROF[name]
@@ -533,7 +533,7 @@ def stepstr(self) -> str:
     @cached_property
     def centers(self) -> ndarray:
         """Radial position of cell centers."""
-        return self._rprofs["r"].values + self.bounds[0]
+        return self._rprofs["r"].to_numpy() + self.bounds[0]
 
     @cached_property
     def walls(self) -> ndarray:

stagpy/commands.py

@@ -71,7 +71,7 @@ def info_cmd() -> None:
                     pandas.Series(data=dimensions, index=series.index, name="dim"),
                 ],
                 axis=1,
-            )
+            )  # type: ignore
         print(indent(series.to_string(header=False), "  "))
         print()
 

stagpy/stagyydata.py

@@ -257,7 +257,7 @@ def _tseries(self) -> DataFrame:
 
     def __getitem__(self, name: str) -> Tseries:
         if name in self._tseries.columns:
-            series = self._tseries[name].values
+            series = self._tseries[name].to_numpy()
             time = self.time
             if name in phyvars.TIME:
                 meta = phyvars.TIME[name]
@@ -308,7 +308,7 @@ def tslice(
     @property
     def time(self) -> ndarray:
         """Time vector."""
-        return self._tseries["t"].values
+        return self._tseries["t"].to_numpy()
 
     @property
     def isteps(self) -> ndarray:
@@ -320,7 +320,7 @@ def isteps(self) -> ndarray:
 
     def at_step(self, istep: int) -> Series:
         """Time series output for a given step."""
-        return self._tseries.loc[istep]
+        return self._tseries.loc[istep]  # type: ignore
 
 
 class _RprofsAveraged(_step._Rprofs):
@@ -816,7 +816,7 @@ def _rprof_and_times(self) -> Tuple[Dict[int, DataFrame], Optional[DataFrame]]:
         return stagyyparsers.rprof(rproffile, list(phyvars.RPROF.keys()))
 
     @property
-    def rtimes(self) -> DataFrame:
+    def rtimes(self) -> Optional[DataFrame]:
         """Radial profiles times."""
         return self._rprof_and_times[1]
 

stagpy/stagyyparsers.py

@@ -80,7 +80,7 @@ def time_series(timefile: Path, colnames: List[str]) -> Optional[DataFrame]:
         memory_map=True,
         on_bad_lines="skip",
     )
-    data = data.apply(pd.to_numeric, raw=True, errors="coerce")
+    data = data.apply(pd.to_numeric, raw=True, errors="coerce")  # type: ignore
 
     # detect useless lines produced when run is restarted
     rows_to_del = []
@@ -97,7 +97,7 @@ def time_series(timefile: Path, colnames: List[str]) -> Optional[DataFrame]:
 
     ncols = data.shape[1]
     _tidy_names(colnames, ncols)
-    data.columns = colnames
+    data.columns = pd.Index(colnames)
 
     return data
 
@@ -126,7 +126,7 @@ def time_series_h5(timefile: Path, colnames: List[str]) -> Optional[DataFrame]:
         h5names = h5f["names"].asstr()[len(colnames) + 1 :]
         _tidy_names(colnames, ncols, h5names)
         data = dset[()]
-    pdf = pd.DataFrame(data[:, 1:], index=np.int_(data[:, 0]), columns=colnames)
+    pdf = pd.DataFrame(data[:, 1:], index=data[:, 0].astype(np.int64), columns=colnames)
     # remove duplicated lines in case of restart
     return pdf.loc[~pdf.index.duplicated(keep="last")]
 
@@ -198,22 +198,22 @@ def rprof(
         memory_map=True,
         on_bad_lines="skip",
     )
-    data = data.apply(pd.to_numeric, raw=True, errors="coerce")
+    data = data.apply(pd.to_numeric, raw=True, errors="coerce")  # type: ignore
 
     isteps = _extract_rsnap_isteps(rproffile, data)
 
-    data = {}
+    all_data = {}
     for istep, _, step_df in isteps:
-        step_df.index = range(step_df.shape[0])  # check whether necessary
+        step_df.index = pd.RangeIndex(step_df.shape[0])  # check whether necessary
         step_cols = list(colnames)
         _tidy_names(step_cols, step_df.shape[1])
-        step_df.columns = step_cols
-        data[istep] = step_df
+        step_df.columns = pd.Index(step_cols)
+        all_data[istep] = step_df
 
     df_times = pd.DataFrame(
-        list(map(itemgetter(1), isteps)), index=map(itemgetter(0), isteps)
+        list(map(itemgetter(1), isteps)), index=pd.Index(map(itemgetter(0), isteps))
     )
-    return data, df_times
+    return all_data, df_times
 
 
 def rprof_h5(
@@ -250,7 +250,7 @@ def rprof_h5(
             isteps.append((istep, dset.attrs["time"]))
 
     df_times = pd.DataFrame(
-        list(map(itemgetter(1), isteps)), index=map(itemgetter(0), isteps)
+        list(map(itemgetter(1), isteps)), index=pd.Index(map(itemgetter(0), isteps))
     )
     return data, df_times
 
@@ -297,7 +297,7 @@ def refstate(
         memory_map=True,
         on_bad_lines="skip",
     )
-    data = data.apply(pd.to_numeric, raw=True, errors="coerce")
+    data = data.apply(pd.to_numeric, raw=True, errors="coerce")  # type: ignore
     # drop lines corresponding to metadata
     data.dropna(subset=[0], inplace=True)
     isystem = -1

tox.ini

@@ -22,6 +22,7 @@ setenv = STAGPY_ISOLATED=True
 [testenv:mypy]
 deps =
     mypy>=1.9.0
+    pandas-stubs>=2.0.3.230814
 commands=
     mypy --install-types --non-interactive stagpy/ tests/