fix mypy errors

ahcore/backends.py

@@ -2,7 +2,7 @@
 
 import pyvips
 from dlup.backends.common import AbstractSlideBackend
-from dlup.types import PathLike
+from dlup.types import PathLike  # type: ignore
 
 from ahcore.readers import StitchingMode, ZarrFileImageReader
 
@@ -14,7 +14,7 @@ def __init__(self, filename: PathLike, stitching_mode: StitchingMode | str = Sti
         self._spacings = [(self._reader.mpp, self._reader.mpp)]
 
     @property
-    def size(self):
+    def size(self) -> tuple[int, int]:
         return self._reader.size
 
     @property
@@ -34,11 +34,11 @@ def properties(self) -> dict[str, Any]:
         return self._reader.metadata
 
     @property
-    def magnification(self):
+    def magnification(self) -> float | None:
         return None
 
     def read_region(self, coordinates: tuple[int, int], level: int, size: tuple[int, int]) -> pyvips.Image:
         return self._reader.read_region(coordinates, level, size)
 
-    def close(self):
+    def close(self) -> None:
         self._reader.close()

ahcore/cli/tiling.py

@@ -20,7 +20,7 @@
 import PIL.Image
 import PIL.ImageCms
 from dlup import SlideImage
-from dlup.backends import ImageBackend
+from dlup.backends import ImageBackend  # type: ignore
 from dlup.data.dataset import TiledWsiDataset
 from dlup.tiling import GridOrder, TilingMode
 from PIL import Image

ahcore/data/dataset.py

@@ -12,7 +12,7 @@
 import pytorch_lightning as pl
 import torch
 from dlup.data.dataset import Dataset, TiledWsiDataset
-from torch.utils.data import DataLoader, DistributedSampler, Sampler
+from torch.utils.data import DataLoader, DistributedSampler, RandomSampler, Sampler, SequentialSampler
 
 from ahcore.utils.data import DataDescription, basemodel_to_uuid
 from ahcore.utils.debug_utils import time_it
@@ -226,7 +226,6 @@ def dataset_iterator() -> Generator[TiledWsiDataset, None, None]:
 
         setattr(self, f"_{stage}_data_iterator", dataset_iterator())
 
-    @time_it
     def _construct_concatenated_dataloader(
         self, data_iterator: Iterator[_DlupDataset], batch_size: int, stage: str, distributed: bool = False
     ) -> Optional[DataLoader[DlupDatasetSample]]:
@@ -260,7 +259,7 @@ def construct_dataset() -> ConcatDataset:
 
         sampler: Sampler[int]
         if stage == "fit":
-            sampler = torch.utils.data.RandomSampler(data_source=dataset)
+            sampler = RandomSampler(data_source=dataset)
 
         elif stage == "predict" and distributed:
             # this is necessary because Lightning changes backend logic for predict
@@ -271,7 +270,7 @@ def construct_dataset() -> ConcatDataset:
             )
 
         else:
-            sampler = torch.utils.data.SequentialSampler(data_source=dataset)
+            sampler = SequentialSampler(data_source=dataset)
 
         return DataLoader(
             dataset,

ahcore/lit_module.py

@@ -10,7 +10,7 @@
 from typing import Any
 
 import pytorch_lightning as pl
-import torch.optim.optimizer
+import torch
 from pytorch_lightning.trainer.states import TrainerFn
 from torch import nn
 
@@ -39,7 +39,7 @@ class AhCoreLightningModule(pl.LightningModule):
     def __init__(
         self,
         model: nn.Module | BaseAhcoreJitModel,
-        optimizer: torch.optim.Optimizer,  # noqa
+        optimizer: torch.optim.optimizer.Optimizer,  # noqa
         data_description: DataDescription,
         loss: nn.Module | None = None,
         augmentations: dict[str, nn.Module] | None = None,

ahcore/utils/callbacks.py

@@ -28,7 +28,9 @@
 logging.getLogger("pyvips").setLevel(logging.ERROR)
 
 
-def _validate_annotations(data_description, annotations: Optional[WsiAnnotations]) -> Optional[WsiAnnotations]:
+def _validate_annotations(
+    data_description: Optional[DataDescription], annotations: Optional[WsiAnnotations]
+) -> Optional[WsiAnnotations]:
     if annotations is None:
         return None
 
@@ -146,7 +148,7 @@ def __getitem__(self, idx: int) -> dict[str, Any]:
             if roi is not None:
                 sample["roi"] = roi.astype(np.uint8)
             else:
-                sample["roi"] = None  # type: ignore
+                sample["roi"] = None
             sample["target"] = target
 
         return sample

ahcore/utils/io.py

@@ -276,7 +276,7 @@ def validate_checkpoint_paths(config: DictConfig) -> DictConfig:
         return config
 
 
-def get_git_hash():
+def get_git_hash() -> Optional[str]:
     try:
         # Check if we're in a git repository
         subprocess.run(["git", "rev-parse", "--is-inside-work-tree"], check=True, capture_output=True, text=True)

ahcore/utils/manifest.py

@@ -12,7 +12,7 @@
 
 from dlup import SlideImage
 from dlup.annotations import WsiAnnotations
-from dlup.backends import ImageBackend
+from dlup.backends import ImageBackend  # type: ignore
 from dlup.data.dataset import RegionFromWsiDatasetSample, TiledWsiDataset, TileSample
 from dlup.tiling import GridOrder, TilingMode
 from pydantic import BaseModel

ahcore/writers.py

@@ -68,7 +68,7 @@ class WriterMetadata(NamedTuple):
     num_channels: int
     dtype: str
     grid_offset: tuple[int, int] | None
-    versions: dict[str, str]
+    versions: dict[str, str | None]
 
 
 class Writer(abc.ABC):