fix types as per pyright

.pre-commit-config.yaml

@@ -8,3 +8,12 @@ repos:
         args: [--fix]
       # Run the formatter.
       - id: ruff-format
+
+  - repo: local
+    hooks:
+      - id: pyright-verifytypes
+        name: pyright verifytypes
+        entry: bash -c 'pip install . >/dev/null && pyright --verifytypes swvo --ignoreexternal'
+        language: system
+        types: [python]
+        pass_filenames: false

pyproject.toml

@@ -42,6 +42,9 @@ Tracker = "https://github.com/GFZ/SWVO/issues"
 [tool.setuptools.packages.find]
 include = ["swvo*"]
 
+[tool.setuptools.package-data]
+"swvo" = ["py.typed"]
+
 
 [tool.coverage.run]
 source = ["swvo"]

requirements.txt

@@ -16,4 +16,4 @@ tqdm
 pytest-mock
 netcdf4
 icecream
-
+pyright

ruff.toml

@@ -79,3 +79,5 @@ docstring-code-format = false
 # This only has an effect when the `docstring-code-format` setting is
 # enabled.
 docstring-code-line-length = "dynamic"
+
+

swvo/io/RBMDataSet/RBMDataSet.py

@@ -182,16 +182,16 @@ def __init__(
             self._file_loading_mode = True
             self._enable_dict_loading = enable_dict_loading
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self._satellite}, {self._instrument}, {self._mfm})"
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.__repr__()
 
-    def __dir__(self):
+    def __dir__(self) -> list[str]:
         return list(super().__dir__()) + [var.var_name for var in VariableEnum]
 
-    def __getattr__(self, name: str):
+    def __getattr__(self, name: str) -> NDArray[np.float64]:
         # Avoid recursion for internal attributes
         if name.startswith("_"):
             raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
@@ -303,7 +303,7 @@ def update_from_dict(self, source_dict: dict[str, VariableLiteral]) -> RBMDataSe
                 raise VariableNotFoundError(msg)
         return self
 
-    def get_var(self, var: VariableEnum):
+    def get_var(self, var: VariableEnum) -> NDArray[np.float64]:
         return getattr(self, var.var_name)
 
     def _create_date_list(self) -> list[dt.datetime]:
@@ -349,15 +349,15 @@ def get_satellite_name(self) -> str:
     def get_satellite_and_instrument_name(self) -> str:
         return self._satellite.sat_name + "_" + self._instrument.instrument_name
 
-    def set_file_path_stem(self, file_path_stem: Path):
+    def set_file_path_stem(self, file_path_stem: Path) -> RBMDataSet:
         self._file_path_stem = file_path_stem
         return self
 
-    def set_file_name_stem(self, file_name_stem: Path):
+    def set_file_name_stem(self, file_name_stem: Path) -> RBMDataSet:
         self._file_path_stem = file_name_stem
         return self
 
-    def set_file_cadence(self, file_cadence: FileCadenceEnum):
+    def set_file_cadence(self, file_cadence: FileCadenceEnum) -> RBMDataSet:
         self._file_cadence = file_cadence
         self._date_of_files = self._create_date_list()
         return self

swvo/io/RBMDataSet/RBMNcDataSet.py

@@ -13,15 +13,18 @@
 from numpy.typing import NDArray
 
 from swvo.io.RBMDataSet import (
+    RBMDataSet,
+)
+from swvo.io.RBMDataSet.custom_enums import (
     FolderTypeEnum,
     InstrumentLike,
+    MfmEnumLiteral,
     MfmLike,
-    RBMDataSet,
     SatelliteLike,
     Variable,
     VariableEnum,
+    VariableLiteral,
 )
-from swvo.io.RBMDataSet.custom_enums import MfmEnumLiteral, VariableLiteral
 from swvo.io.RBMDataSet.utils import join_var
 
 

swvo/io/RBMDataSet/bin_and_interpolate_to_model_grid.py

@@ -10,20 +10,19 @@
 from functools import partial
 from multiprocessing import Pool
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal
+from typing import Literal
 
 import numpy as np
 from icecream import ic
 from matplotlib import pyplot as plt
 from numpy.typing import NDArray
 from tqdm import tqdm
 
-if TYPE_CHECKING:
-    from swvo.io.RBMDataSet import RBMDataSet, RBMDataSetElPaso
+from swvo.io.RBMDataSet import RBMDataSet
 
 
 def bin_and_interpolate_to_model_grid(
-    self: RBMDataSet | RBMDataSetElPaso,
+    self: RBMDataSet,
     sim_time: list[datetime],
     grid_R: NDArray[np.float64],
     grid_mu_V: NDArray[np.float64],
@@ -383,7 +382,7 @@ def plot_debug_figures_plasmasphere(
     grid_P: NDArray[np.float64] | None,
     grid_R: NDArray[np.float64],
     debug_plot_settings: DebugPlotSettings,
-):
+) -> None:
     print("\tPlot debug features...")
 
     dt = sim_time[1] - sim_time[0]
@@ -462,7 +461,7 @@ def plot_debug_figures(
     grid_K: NDArray[np.float64],
     mu_or_V: Literal["Mu", "V"],
     debug_plot_settings: DebugPlotSettings,
-):
+) -> None:
     print("\tPlot debug features...")
 
     dt = sim_time[1] - sim_time[0]

swvo/io/RBMDataSet/interp_functions.py

@@ -10,14 +10,12 @@
 from enum import Enum
 from functools import partial
 from multiprocessing import Pool
-from typing import TYPE_CHECKING
 
 import numpy as np
 from numpy.typing import NDArray
 from tqdm import tqdm
 
-if TYPE_CHECKING:
-    from swvo.io.RBMDataSet import RBMDataSet
+from swvo.io.RBMDataSet import RBMDataSet
 
 
 class TargetType(Enum):
@@ -222,13 +220,13 @@ def _interp_psd_parallel(
 
 
 def interp_psd(
-    self,
+    self: RBMDataSet,
     target_K: list[float],
     target_type: TargetType,
     target_mu: list[float] = None,
     target_V: list[float] = None,
-    n_threads=10,
-):
+    n_threads: int=10,
+) -> NDArray[np.float64]:
     if not isinstance(target_K, Iterable):
         target_K = [target_K]
 

swvo/io/RBMDataSet/utils.py

@@ -102,7 +102,7 @@ def python2matlab(datenum: datetime) -> float:
     return mdn.toordinal() + round(frac, 6)
 
 
-def matlab2python(datenum: float | Iterable) -> Iterable[datetime] | datetime:
+def matlab2python(datenum: float | Iterable[float]) -> Iterable[datetime] | datetime:
     """Convert MATLAB datenum to Python datetime."""
     warnings.filterwarnings("ignore", message="Discarding nonzero nanoseconds in conversion")
 

swvo/io/dst/wdc.py

@@ -53,7 +53,7 @@ def __init__(self, data_dir: Optional[Path] = None) -> None:
 
             data_dir = os.environ.get(self.ENV_VAR_NAME)
 
-        self.data_dir = Path(data_dir)
+        self.data_dir: Path = Path(data_dir)
         self.data_dir.mkdir(parents=True, exist_ok=True)
 
         logging.info(f"WDC Dst data directory: {self.data_dir}")

swvo/io/f10_7/swpc.py

@@ -55,7 +55,7 @@ def __init__(self, data_dir: Optional[Path] = None) -> None:
                 raise ValueError(msg)
             data_dir = os.environ.get(self.ENV_VAR_NAME)
 
-        self.data_dir = Path(data_dir)
+        self.data_dir: Path = Path(data_dir)
         self.data_dir.mkdir(parents=True, exist_ok=True)
 
         logging.info(f"SWPC F10.7 data directory: {self.data_dir}")

swvo/io/hp/ensemble.py

@@ -56,15 +56,15 @@ def __init__(self, index: str, data_dir: Optional[Path] = None) -> None:
 
             data_dir = os.environ.get(self.ENV_VAR_NAME)
 
-        self.data_dir = Path(data_dir)
+        self.data_dir: Path = Path(data_dir)
 
         logging.info(f"{self.index.upper()} Ensemble data directory: {self.data_dir}")
 
         if not self.data_dir.exists():
             msg = f"Data directory {self.data_dir} does not exist! Impossible to retrive data!"
             raise FileNotFoundError(msg)
 
-        self.index_number = index[2:]
+        self.index_number: int = index[2:]
 
     def read(self, start_time: datetime, end_time: datetime) -> list[pd.DataFrame]:
         """

swvo/io/hp/gfz.py

@@ -56,9 +56,9 @@ def __init__(self, index: str, data_dir: Optional[Path] = None) -> None:
 
             data_dir = os.environ.get(self.ENV_VAR_NAME)
 
-        self.data_dir = Path(data_dir)
+        self.data_dir: Path = Path(data_dir)
         self.data_dir.mkdir(parents=True, exist_ok=True)
-        self.index_number = index[2:]
+        self.index_number: int = index[2:]
 
         logging.info(f"{self.index.upper()} GFZ data directory: {self.data_dir}")
 

swvo/io/kp/ensemble.py

@@ -49,14 +49,14 @@ def __init__(self, data_dir: Optional[Path] = None) -> None:
 
             data_dir = os.environ.get(self.ENV_VAR_NAME)
 
-        self.data_dir = Path(data_dir)
+        self.data_dir: Path = Path(data_dir)
 
         logging.info(f"Kp Ensemble data directory: {self.data_dir}")
 
         if not self.data_dir.exists():
             raise FileNotFoundError(f"Data directory {self.data_dir} does not exist! Impossible to retrive data!")
 
-    def read(self, start_time: datetime, end_time: datetime) -> list:
+    def read(self, start_time: datetime, end_time: datetime) -> list[pd.DataFrame]:
         """Read Kp ensemble data for the requested period.
 
         Parameters

swvo/io/kp/niemegk.py

@@ -55,7 +55,7 @@ def __init__(self, data_dir: Optional[Path] = None) -> None:
 
             data_dir = os.environ.get(self.ENV_VAR_NAME)
 
-        self.data_dir = Path(data_dir)
+        self.data_dir: Path = Path(data_dir)
         self.data_dir.mkdir(parents=True, exist_ok=True)
 
         logging.info(f"Kp Niemegk  data directory: {self.data_dir}")