Adopt ty and refactor core typing everywhere (#1414)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: fepegar

.pre-commit-config.yaml

@@ -37,6 +37,18 @@ repos:
       - id: pyupgrade
         args: ['--py310-plus']
 
+  - repo: local
+    hooks:
+      - id: ty-check
+        name: ty check
+        entry: uv run --group types ty check --output-format concise src
+        language: system
+        always_run: true
+        pass_filenames: false
+        require_serial: true
+        types_or: [python, pyi]
+
 ci:
   autoupdate_commit_msg: Autoupdate pre-commit hooks
   autoupdate_schedule: quarterly
+  skip: [ty-check]

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
     "deprecated>=1.2",
     "einops>=0.3",
     "humanize>=0.1",
+    "jaxtyping>=0.2",
     "nibabel>=3",
     "numpy>=1.20",
     "packaging>=20",
@@ -85,7 +86,8 @@ doc = [
 maintain = ["bump-my-version"]
 quality = ["ruff"]
 types = [
-    "mypy",
+    "matplotlib",
+    "ty",
     "pandas-stubs",
     "pillow",
     "pip",
@@ -117,29 +119,6 @@ filename = "pyproject.toml"
 search = 'version = "{current_version}"'
 replace = 'version = "{new_version}"'
 
-[tool.mypy]
-pretty = true
-
-[[tool.mypy.overrides]]
-module = "duecredit.*"
-ignore_missing_imports = true
-
-[[tool.mypy.overrides]]
-module = "huggingface_hub.*"
-ignore_missing_imports = true
-
-[[tool.mypy.overrides]]
-module = "matplotlib.*"
-ignore_missing_imports = true
-
-[[tool.mypy.overrides]]
-module = "nibabel.*"
-ignore_missing_imports = true
-
-[[tool.mypy.overrides]]
-module = "scipy.*"
-ignore_missing_imports = true
-
 [tool.pytest.ini_options]
 addopts = ["-ra", "--strict-config", "--strict-markers"]
 filterwarnings = [

src/torchio/data/dataset.py

@@ -11,7 +11,7 @@
 from .subject import Subject
 
 
-class SubjectsDataset(Dataset):
+class SubjectsDataset(Dataset[Subject]):
     """Base TorchIO dataset.
 
     Reader of 3D medical images that directly inherits from the PyTorch
@@ -61,12 +61,12 @@ class SubjectsDataset(Dataset):
     def __init__(
         self,
         subjects: Sequence[Subject],
-        transform: Callable | None = None,
+        transform: Callable[[Subject], Subject] | None = None,
         load_getitem: bool = True,
     ):
         self._parse_subjects_list(subjects)
         self._subjects = subjects
-        self._transform: Callable | None
+        self._transform: Callable[[Subject], Subject] | None
         self.set_transform(transform)
         self.load_getitem = load_getitem
 
@@ -104,7 +104,7 @@ def from_batch(cls, batch: dict) -> SubjectsDataset:
         subjects: list[Subject] = get_subjects_from_batch(batch)
         return cls(subjects)
 
-    def dry_iter(self):
+    def dry_iter(self) -> Sequence[Subject]:
         """Return the internal list of subjects.
 
         This can be used to iterate over the subjects without loading the data
@@ -115,7 +115,10 @@ def dry_iter(self):
         """
         return self._subjects
 
-    def set_transform(self, transform: Callable | None) -> None:
+    def set_transform(
+        self,
+        transform: Callable[[Subject], Subject] | None,
+    ) -> None:
         """Set the `transform` attribute.
 
         Args: