Apply suggestions from code review

Co-authored-by: Chris Markiewicz <[email protected]>

Author: oesteban

src/nifreeze/estimator.py

@@ -26,7 +26,7 @@
 
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Self
+from typing import Self, TypeVar
 
 from tqdm import tqdm
 
@@ -39,10 +39,13 @@
 from nifreeze.utils import iterators
 
 
+DatasetT = TypeVar("DatasetT", bound=BaseDataset)
+
+
 class Filter:
     """Alters an input data object (e.g., downsampling)."""
 
-    def run(self, dataset: BaseDataset, **kwargs):
+    def run(self, dataset: DatasetT, **kwargs) -> DatasetT:
         """
         Trigger execution of the designated filter.
 
@@ -53,8 +56,8 @@ def run(self, dataset: BaseDataset, **kwargs):
 
         Returns
         -------
-        :obj:`~nifreeze.estimator.Estimator`
-            The estimator, after fitting.
+        dataset : :obj:`~nifreeze.data.base.BaseDataset`
+            The dataset, after filtering.
 
         """
         return dataset
@@ -69,7 +72,7 @@ def __init__(
         self,
         model: BaseModel | str,
         strategy: str = "random",
-        prev: Self | None = None,
+        prev: Estimator | Filter | None = None,
         model_kwargs: dict | None = None,
         **kwargs,
     ):
@@ -79,7 +82,7 @@ def __init__(
         self._model_kwargs = model_kwargs or {}
         self._align_kwargs = kwargs or {}
 
-    def run(self, dataset: BaseDataset, **kwargs):
+    def run(self, dataset: DatasetT, **kwargs) -> Self:
         """
         Trigger execution of the workflow this estimator belongs.
 

src/nifreeze/model/base.py

@@ -22,6 +22,7 @@
 #
 """Base infrastructure for nifreeze's models."""
 
+from abc import abstractmethod
 from warnings import warn
 
 import numpy as np
@@ -95,7 +96,8 @@ def __init__(self, dataset, **kwargs):
         if dataset.brainmask is None:
             warn(mask_absence_warn_msg, stacklevel=2)
 
-    def fit_predict(self, *_, **kwargs):
+    @abstractmethod
+    def fit_predict(self, index, **kwargs) -> np.ndarray:
         """Fit and predict the indicate index of the dataset (abstract signature)."""
         raise NotImplementedError("Cannot call fit_predict() on a BaseModel instance.")
 
@@ -139,7 +141,7 @@ def __init__(self, dataset, stat="median", **kwargs):
         super().__init__(dataset, **kwargs)
         self._stat = stat
 
-    def fit_predict(self, index, *_, **kwargs):
+    def fit_predict(self, index, **kwargs):
         """
         Return the expectation map.