junifer 0.0.7.dev332__py3-none-any.whl → 0.0.8__py3-none-any.whl

junifer/datagrabber/aomic/tests/test_piop2.py

@@ -7,14 +7,41 @@
 #          Synchon Mandal <s.mandal@fz-juelich.de>
 # License: AGPL
 
-from typing import Optional, Union
+from collections.abc import Iterator
+from pathlib import Path
 
 import pytest
+from pydantic import AnyUrl
 
 from junifer.datagrabber import DataladAOMICPIOP2
+from junifer.testing import config_override
 
 
-URI = "https://gin.g-node.org/juaml/datalad-example-aomicpiop2"
+URI = AnyUrl(
+    "https://cerebra.fz-juelich.de/junifer/datalad-example-aomicpiop2"
+)
+
+
+@pytest.fixture(scope="module")
+def aomic_datadir(tmp_path_factory) -> Iterator[Path]:
+    """Return the path to the AOMIC data directory.
+
+    Parameters
+    ----------
+    tmp_path_factory : pathlib.Path
+        The path to the test directory.
+
+    Returns
+    -------
+    pathlib.Path
+        The path to the AOMIC data directory.
+
+    """
+    datadir = tmp_path_factory.mktemp("aomic_data")
+    datadir.mkdir(parents=True, exist_ok=True)
+    dg = DataladAOMICPIOP2(uri=URI, types=["T1w"], datadir=datadir)
+    with config_override("datagrabber.skipidcheck", True), dg:
+        yield datadir
 
 
 @pytest.mark.parametrize(
@@ -23,18 +50,11 @@ URI = "https://gin.g-node.org/juaml/datalad-example-aomicpiop2"
         (
             "BOLD",
             ["confounds", "mask", "reference"],
-            None,
+            "restingstate",
             "MNI152NLin2009cAsym",
         ),
-        ("BOLD", ["confounds", "mask", "reference"], None, "native"),
         (
-            "BOLD",
-            ["confounds", "mask", "reference"],
-            ["restingstate"],
-            "MNI152NLin2009cAsym",
-        ),
-        (
-            "BOLD",
+            ["BOLD"],
             ["confounds", "mask", "reference"],
             ["restingstate", "stopsignal"],
             "MNI152NLin2009cAsym",
@@ -46,72 +66,79 @@ URI = "https://gin.g-node.org/juaml/datalad-example-aomicpiop2"
             "MNI152NLin2009cAsym",
         ),
         (
-            "BOLD",
+            ["BOLD"],
             ["confounds", "mask", "reference"],
-            ["workingmemory"],
+            "workingmemory",
             "MNI152NLin2009cAsym",
         ),
-        ("T1w", ["mask"], None, "MNI152NLin2009cAsym"),
-        ("T1w", ["mask"], None, "native"),
-        ("VBM_CSF", None, None, "MNI152NLin2009cAsym"),
-        ("VBM_CSF", None, None, "native"),
-        ("VBM_GM", None, None, "MNI152NLin2009cAsym"),
-        ("VBM_GM", None, None, "native"),
-        ("VBM_WM", None, None, "MNI152NLin2009cAsym"),
-        ("VBM_WM", None, None, "native"),
-        ("DWI", None, None, "MNI152NLin2009cAsym"),
-        ("FreeSurfer", None, None, "MNI152NLin2009cAsym"),
+        (["T1w"], ["mask"], "restingstate", "MNI152NLin2009cAsym"),
+        ("T1w", ["mask"], ["restingstate"], "native"),
+        (["VBM_CSF"], None, "restingstate", "MNI152NLin2009cAsym"),
+        ("VBM_CSF", None, ["restingstate"], "native"),
+        (["VBM_GM"], None, "restingstate", "MNI152NLin2009cAsym"),
+        ("VBM_GM", None, ["restingstate"], "native"),
+        (["VBM_WM"], None, "restingstate", "MNI152NLin2009cAsym"),
+        ("VBM_WM", None, ["restingstate"], "native"),
+        (["DWI"], None, "restingstate", "MNI152NLin2009cAsym"),
+        (["FreeSurfer"], None, ["restingstate"], "MNI152NLin2009cAsym"),
     ],
 )
 def test_DataladAOMICPIOP2(
-    type_: str,
-    nested_types: Optional[list[str]],
-    tasks: Optional[list[str]],
+    type_: str | list[str],
+    nested_types: list[str] | None,
+    tasks: str | list[str],
     space: str,
+    aomic_datadir: Path,
 ) -> None:
     """Test DataladAOMICPIOP2 DataGrabber.
 
     Parameters
     ----------
-    type_ : str
+    type_ : str or list of str
         The parametrized type.
     nested_types : list of str or None
         The parametrized nested types.
-    tasks : list of str or None
+    tasks : str or list of str
         The parametrized task values.
     space: str
         The parametrized space.
+    aomic_datadir: Path
+        The path to the AOMIC data directory.
 
     """
-    dg = DataladAOMICPIOP2(types=type_, tasks=tasks, space=space)
-    # Set URI to Gin
-    dg.uri = URI
-
+    dg = DataladAOMICPIOP2(
+        uri=URI, types=type_, tasks=tasks, space=space, datadir=aomic_datadir
+    )
     with dg:
-        # Get all elements
         all_elements = dg.get_elements()
-        # Get test element
         test_element = all_elements[0]
-        # Get test element data
         out = dg[test_element]
+        # Elements only have the subject if no data type of the tasks is
+        # grabbed
+        if not isinstance(test_element, tuple):
+            test_element = (test_element,)
         # Assert data type
-        assert type_ in out
-        # Check task name if BOLD
-        if type_ == "BOLD" and tasks is not None:
-            assert test_element[1] in out[type_]["path"].name
-        assert out[type_]["path"].exists()
-        assert out[type_]["path"].is_file()
-        # Asserts data type metadata
-        assert "meta" in out[type_]
-        meta = out[type_]["meta"]
-        assert "element" in meta
-        assert "subject" in meta["element"]
-        assert test_element[0] == meta["element"]["subject"]
-        # Assert nested data type if not None
-        if nested_types is not None:
-            for nested_type in nested_types:
-                assert out[type_][nested_type]["path"].exists()
-                assert out[type_][nested_type]["path"].is_file()
+        if isinstance(type_, str):
+            type_ = [type_]
+        for t in type_:
+            assert t in out
+            # Check task name if BOLD
+            if t == "BOLD":
+                assert test_element[1] in out[t]["path"].name
+            assert out[t]["path"].exists()
+            assert out[t]["path"].is_file()
+            # Asserts data type metadata
+            assert "meta" in out[t]
+            meta = out[t]["meta"]
+            assert "element" in meta
+            assert "subject" in meta["element"]
+            assert test_element[0] == meta["element"]["subject"]
+            assert list(meta["element"]) == dg.get_element_keys()
+            # Assert nested data type if not None
+            if nested_types is not None:
+                for nested_type in nested_types:
+                    assert out[t][nested_type]["path"].exists()
+                    assert out[t][nested_type]["path"].is_file()
 
 
 @pytest.mark.parametrize(
@@ -130,7 +157,8 @@ def test_DataladAOMICPIOP2(
     ],
 )
 def test_DataladAOMICPIOP2_partial_data_access(
-    types: Union[str, list[str]],
+    types: str | list[str],
+    aomic_datadir: Path,
 ) -> None:
     """Test DataladAOMICPIOP2 DataGrabber partial data access.
 
@@ -138,42 +166,17 @@ def test_DataladAOMICPIOP2_partial_data_access(
     ----------
     types : str or list of str
         The parametrized types.
+    aomic_datadir : pathlib.Path
+        The path to the AOMIC data directory.
 
     """
-    dg = DataladAOMICPIOP2(types=types)
-    # Set URI to Gin
-    dg.uri = URI
-
+    dg = DataladAOMICPIOP2(uri=URI, types=types, datadir=aomic_datadir)
     with dg:
-        # Get all elements
         all_elements = dg.get_elements()
-        # Get test element
         test_element = all_elements[0]
-        # Get test element data
         out = dg[test_element]
         # Assert data type
-        if isinstance(types, list):
-            for type_ in types:
-                assert type_ in out
-        else:
-            assert types in out
-
-
-def test_DataladAOMICPIOP2_incorrect_data_type() -> None:
-    """Test DataladAOMICPIOP2 DataGrabber incorrect data type."""
-    with pytest.raises(
-        ValueError, match="`patterns` must contain all `types`"
-    ):
-        _ = DataladAOMICPIOP2(types="Vesta")
-
-
-def test_DataladAOMICPIOP2_invalid_tasks():
-    """Test DataladAOMICIDPIOP2 DataGrabber invalid tasks."""
-    with pytest.raises(
-        ValueError,
-        match=(
-            "thisisnotarealtask is not a valid task in "
-            "the AOMIC PIOP2 dataset!"
-        ),
-    ):
-        DataladAOMICPIOP2(tasks="thisisnotarealtask")
+        if isinstance(types, str):
+            types = [types]
+        for t_ in types:
+            assert t_ in out

junifer/datagrabber/base.py

@@ -7,57 +7,140 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
+from enum import Enum
 from pathlib import Path
-from typing import Union
+from typing import Annotated, Any
+
+import structlog
+from aenum import Enum as AEnum
+from pydantic import BaseModel, BeforeValidator, ConfigDict
 
 from ..pipeline import UpdateMetaMixin
 from ..typing import Element, Elements
-from ..utils import logger, raise_error
+from ..utils import ensure_list, raise_error
 
 
-__all__ = ["BaseDataGrabber"]
+__all__ = ["BaseDataGrabber", "DataType"]
 
+_log = structlog.get_logger("junifer")
+logger = _log.bind(pkg="datagrabber", step="datagrabber")
 
-class BaseDataGrabber(ABC, UpdateMetaMixin):
-    """Abstract base class for DataGrabber.
 
-    For every interface that is required, one needs to provide a concrete
-    implementation of this abstract class.
+class DataType(str, AEnum):
+    """Accepted data type."""
+
+    T1w = "T1w"
+    T2w = "T2w"
+    BOLD = "BOLD"
+    Warp = "Warp"
+    VBM_GM = "VBM_GM"
+    VBM_WM = "VBM_WM"
+    VBM_CSF = "VBM_CSF"
+    FALFF = "fALFF"
+    GCOR = "GCOR"
+    LCOR = "LCOR"
+    DWI = "DWI"
+    FreeSurfer = "FreeSurfer"
+
+
+def _matches(element: Element, selector: Element) -> bool:
+    """Check whether an element matches a selector.
+
+    Parameters
+    ----------
+    element : ``Element``
+        The element.
+    selector : ``Element``
+        The partial or complete element selector: a value (the element must
+        have it) or a tuple of values (the element must have all of them).
+
+    Returns
+    -------
+    bool
+        Whether the element matches the selector.
+
+    """
+    if not isinstance(element, tuple):
+        element = (element,)
+    if isinstance(selector, tuple):
+        return set(selector).issubset(element)
+    return selector in element
+
+
+def _selectors_as_tuples(selection: Elements) -> list[tuple]:
+    """Get the element selectors as tuples, checking they have the same size.
 
     Parameters
     ----------
-    types : list of str
-        The types of data to be grabbed.
-    datadir : str or pathlib.Path
-        The directory where the data is / will be stored.
+    selection : ``Elements``
+        The list of partial or complete element selectors.
+
+    Returns
+    -------
+    list of tuple
+        The selectors, as tuples (a single value is a tuple of one value).
 
     Raises
     ------
-    TypeError
-        If ``types`` is not a list or if the values are not string.
+    ValueError
+        If the selectors have different numbers of values.
 
     """
+    selectors = [s if isinstance(s, tuple) else (s,) for s in selection]
+    if len({len(selector) for selector in selectors}) > 1:
+        raise_error(
+            msg=(
+                "The element selectors must have the same number of values: "
+                f"{selection}"
+            ),
+            klass=ValueError,
+        )
+    return selectors
 
-    def __init__(self, types: list[str], datadir: Union[str, Path]) -> None:
-        # Validate types
-        if not isinstance(types, list):
-            raise_error(msg="`types` must be a list", klass=TypeError)
-        if any(not isinstance(x, str) for x in types):
-            raise_error(
-                msg="`types` must be a list of strings", klass=TypeError
-            )
-        self.types = types
 
-        # Convert str to Path
-        if not isinstance(datadir, Path):
-            datadir = Path(datadir)
-        self._datadir = datadir
+class BaseDataGrabber(BaseModel, ABC, UpdateMetaMixin):
+    """Abstract base class for data fetcher.
+
+    For every datagrabber, one needs to provide a concrete
+    implementation of this abstract class.
+
+    Parameters
+    ----------
+    types : :enum:`.DataType` or list of variants
+        The data type(s) to grab.
+    datadir : pathlib.Path
+        The path where the data is or will be stored.
+
+    """
 
-        logger.debug("Initializing BaseDataGrabber")
-        logger.debug(f"\t_datadir = {datadir}")
-        logger.debug(f"\ttypes = {types}")
+    model_config = ConfigDict(extra="forbid", use_enum_values=True)
+
+    types: Annotated[
+        DataType | list[DataType],
+        BeforeValidator(ensure_list),
+    ]
+    datadir: Path
+
+    def model_post_init(self, context: Any):  # noqa: D102
+        # Run extra validation for datagrabbers and fail early if needed
+        self.validate_datagrabber_params()
+        # Convert to correct data type
+        self.types = [
+            DataType(t) if isinstance(t, str) else t for t in self.types
+        ]
+        logger.info(
+            f"Parameters: {self.model_dump(mode='json')}",
+            component=self.__class__.__name__,
+        )
+
+    def validate_datagrabber_params(self) -> None:
+        """Run extra logical validation for datagrabber.
+
+        Subclasses can override to provide validation.
+        """
+        pass
 
-    def __iter__(self) -> Iterator:
+    def __iter__(self) -> Iterator[Elements]:
         """Enable iterable support.
 
         Yields
@@ -73,7 +156,7 @@ class BaseDataGrabber(ABC, UpdateMetaMixin):
 
         Parameters
         ----------
-        element : str or tuple of str
+        element : `Element`
             The element to be indexed.
 
         Returns
@@ -83,28 +166,66 @@ class BaseDataGrabber(ABC, UpdateMetaMixin):
             specified element.
 
         """
+        # Convert element to tuple if not already and extract enum values if
+        # present
+        element = (
+            (element,)
+            if not isinstance(element, tuple)
+            else tuple(i.value if isinstance(i, Enum) else i for i in element)
+        )
         logger.info(f"Getting element {element}")
-        # Convert element to tuple if not already
-        if not isinstance(element, tuple):
-            element = (element,)
         # Zip through element keys and actual values to construct element
         # access dictionary
-        named_element: dict = dict(zip(self.get_element_keys(), element))
+        named_element: dict = dict(
+            zip(self.get_element_keys(), element, strict=False)
+        )
         logger.debug(f"Named element: {named_element}")
+        self._check_element(named_element)
         # Fetch element
         out = self.get_item(**named_element)
-        # Update metadata
-        for _, t_val in out.items():
+        # Update metadata: the element being processed and the keys of the
+        # element that the data of each data type depends on
+        for t_type, t_val in out.items():
             self.update_meta(t_val, "datagrabber")
+            element_keys = self.get_type_element_keys(t_type)
             # Conditional for list dtype vals like Warp
-            if isinstance(t_val, list):
-                for entry in t_val:
-                    entry["meta"]["element"] = named_element
-            else:
-                t_val["meta"]["element"] = named_element
+            for entry in t_val if isinstance(t_val, list) else [t_val]:
+                entry["meta"]["element"] = named_element
+                entry["meta"]["_element_keys"] = element_keys
 
         return out
 
+    def get_type_element_keys(self, data_type: str) -> list[str]:
+        """Get the element keys that the data of a data type depends on.
+
+        The element keys (see :meth:`.get_element_keys`) are the keys that
+        the data types to grab depend on.
+
+        Parameters
+        ----------
+        data_type : str
+            The data type.
+
+        Returns
+        -------
+        list of str
+            The element keys, all of them unless a subclass knows which ones
+            the data type depends on (e.g., only the subject for data that is
+            the same for all the tasks of a subject).
+
+        """
+        return self.get_element_keys()
+
+    def _check_element(self, element: dict) -> None:
+        """Check the element can be grabbed.
+
+        Parameters
+        ----------
+        element : dict
+            The element, as given to :meth:`.get_item`.
+
+        """
+
     def __enter__(self) -> "BaseDataGrabber":
         """Context entry."""
         return self
@@ -119,67 +240,120 @@ class BaseDataGrabber(ABC, UpdateMetaMixin):
         Returns
         -------
         list of str
-            The types of data to be grabbed.
+            The data type(s) to grab.
 
         """
-        return self.types.copy()
+        return [x.value if isinstance(x, Enum) else x for x in self.types]
 
     @property
-    def datadir(self) -> Path:
-        """Get data directory path.
+    def fulldir(self) -> Path:
+        """Get complete data directory path.
 
         Returns
         -------
         pathlib.Path
-            Path to the data directory. Can be overridden by subclasses.
+            Complete path to the data directory.
+            Can be overridden by subclasses.
 
         """
-        return self._datadir
+        return self.datadir
 
     def filter(self, selection: Elements) -> Iterator:
         """Filter elements to be grabbed.
 
         Parameters
         ----------
-        selection : list
-            The list of partial or complete element selectors to filter using.
+        selection : ``Elements``
+            The list of partial or complete element selectors to filter using,
+            with the same number of values.
 
         Yields
         ------
         object
             An element that can be indexed by the DataGrabber.
 
+        Raises
+        ------
+        ValueError
+            If the selectors have different numbers of values.
+
         """
+        selectors = _selectors_as_tuples(selection)
+        for element in self.get_elements():
+            if any(_matches(element, selector) for selector in selectors):
+                yield element
+
+    def select_elements(self, selection: Elements) -> list[Element]:
+        """Select the elements to grab.
+
+        All the selectors must have the same number of values (a single value
+        counts as a tuple of one value), so they are either all complete
+        elements or all partial selectors:
+
+        * Complete elements have a value for each element key, in the order of
+          :meth:`.get_element_keys`, e.g., ``("sub-01", "rest")`` for the keys
+          ``["subject", "task"]``. They are the elements to grab, so the
+          elements of the dataset are not needed (they are not searched for).
+          Repeated elements are only returned once.
+        * Partial selectors have fewer values, e.g., ``"sub-01"`` or
+          ``("sub-01",)``. They select the elements of the dataset that match
+          them, i.e., that have all their values for any key (see
+          :meth:`.filter`). The elements of the dataset are searched for once,
+          and each element is returned once, even if several selectors match
+          it: with ``["sub-01", "rest"]``, the element ``("sub-01", "rest")``
+          matches both selectors and is returned once, together with the other
+          tasks of ``sub-01`` and the ``rest`` task of the other subjects.
+
+        Parameters
+        ----------
+        selection : ``Elements``
+            The list of partial or complete element selectors.
+
+        Returns
+        -------
+        list of tuple
+            The selected elements, without duplicates.
+
+        Raises
+        ------
+        ValueError
+            If the selectors have different numbers of values.
+        RuntimeError
+            If a partial selector does not match any element.
 
-        def filter_func(element: Element) -> bool:
-            """Filter element based on selection.
-
-            Parameters
-            ----------
-            element : str or tuple of str
-                The element to be filtered.
-
-            Returns
-            -------
-            bool
-                If the element passes the filter or not.
-
-            """
-            # Convert element to tuple
-            if not isinstance(element, tuple):
-                element = (element,)
-            # Filter based on selection kind
-            if isinstance(selection[0], str):
-                for opt in selection:
-                    if opt in element:
-                        return True
-            elif isinstance(selection[0], tuple):
-                for opt in selection:
-                    if set(opt).issubset(element):
-                        return True
-            return False
-
-        yield from filter(filter_func, self.get_elements())
+        """
+        selectors = _selectors_as_tuples(selection)
+        if not selectors:
+            return []
+        # Complete elements: the selectors are the elements, without the
+        # repeated ones
+        if len(selectors[0]) == len(self.get_element_keys()):
+            # Remove the repeated selectors, keeping the order
+            return list(dict.fromkeys(selectors))
+        # Partial selectors: the elements of the dataset that match them, in a
+        # single pass over the elements. Each element is checked once, so it
+        # is returned once even if several selectors match it
+        selected = []
+        matched = set()
+        for element in self.get_elements():
+            t_matched = [
+                i
+                for i, selector in enumerate(selectors)
+                if _matches(element, selector)
+            ]
+            if t_matched:
+                selected.append(element)
+                matched.update(t_matched)
+        # The selectors that no element matches
+        invalid = [
+            given for i, given in enumerate(selection) if i not in matched
+        ]
+        if invalid:
+            raise_error(
+                msg=f"The following element selectors are invalid:\n{invalid}",
+                klass=RuntimeError,
+            )
+        return selected
 
     @abstractmethod
     def get_element_keys(self) -> list[str]: