Attempting to introduce enough typing in module_registry that mypy in…

… pre-commit github actions will pass. (#201)

pyproject.toml

@@ -154,5 +154,6 @@ module = [
     'astropy.table',
     'GPUtil',
     'toml',
+    'umap',
     ]
 ignore_missing_imports = true

src/fibad/data_sets/__init__.py

@@ -3,4 +3,10 @@
 from .hsc_data_set import HSCDataSet
 from .inference_dataset import InferenceDataSet
 
-__all__ = ["fibad_data_set", "DATA_SET_REGISTRY", "CifarDataSet", "HSCDataSet"]
+__all__ = [
+    "fibad_data_set",
+    "DATA_SET_REGISTRY",
+    "CifarDataSet",
+    "HSCDataSet",
+    "InferenceDataSet",
+]

src/fibad/models/model_registry.py

@@ -1,5 +1,6 @@
 import logging
 from pathlib import Path
+from typing import Any, cast
 
 import torch.nn as nn
 
@@ -27,13 +28,15 @@ def _torch_criterion(self: nn.Module):
     """Load the criterion class using the name defined in the config and
     instantiate it with the arguments defined in the config."""
 
+    config = cast(dict[str, Any], self.config)
+
     # Load the class and get any parameters from the config dictionary
-    criterion_cls = get_or_load_class(self.config["criterion"])
-    criterion_name = self.config["criterion"]["name"]
+    criterion_cls = get_or_load_class(config["criterion"])
+    criterion_name = config["criterion"]["name"]
 
     arguments = {}
-    if criterion_name in self.config:
-        arguments = self.config[criterion_name]
+    if criterion_name in config:
+        arguments = config[criterion_name]
 
     # Print some information about the criterion function and parameters used
     log_string = f"Using criterion: {criterion_name} "
@@ -50,13 +53,15 @@ def _torch_optimizer(self: nn.Module):
     """Load the optimizer class using the name defined in the config and
     instantiate it with the arguments defined in the config."""
 
+    config = cast(dict[str, Any], self.config)
+
     # Load the class and get any parameters from the config dictionary
-    optimizer_cls = get_or_load_class(self.config["optimizer"])
-    optimizer_name = self.config["optimizer"]["name"]
+    optimizer_cls = get_or_load_class(config["optimizer"])
+    optimizer_name = config["optimizer"]["name"]
 
     arguments = {}
-    if optimizer_name in self.config:
-        arguments = self.config[optimizer_name]
+    if optimizer_name in config:
+        arguments = config[optimizer_name]
 
     # Print some information about the optimizer function and parameters used
     log_string = f"Using optimizer: {optimizer_name} "
@@ -127,7 +132,7 @@ def fetch_model_class(runtime_config: dict) -> type[nn.Module]:
     model_cls = None
 
     try:
-        model_cls = get_or_load_class(model_config, MODEL_REGISTRY)
+        model_cls = cast(type[nn.Module], get_or_load_class(model_config, MODEL_REGISTRY))
     except ValueError as exc:
         raise ValueError("Error fetching model class") from exc