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

junifer/datagrabber/multiple.py

@@ -5,23 +5,73 @@
 #          Synchon Mandal <s.mandal@fz-juelich.de>
 # License: AGPL
 
-from typing import Union
+from collections import defaultdict
+from pathlib import Path
+from typing import Annotated
+
+from pydantic import BeforeValidator, ConfigDict
 
 from ..api.decorators import register_datagrabber
-from ..typing import DataGrabberLike
-from ..utils import deep_update, raise_error
-from .base import BaseDataGrabber
+from ..typing import DataGrabberLike, Element
+from ..utils import deep_update, ensure_list, raise_error
+from .base import BaseDataGrabber, DataType
+from .pattern import PatternDataGrabber, _join_elements
+from .pattern_datalad import PatternDataladDataGrabber
 
 
 __all__ = ["MultipleDataGrabber"]
 
 
+def _nested_types_only(
+    datagrabber: DataGrabberLike, data_type: str
+) -> list[str] | None:
+    """Get the nested data types grabbed without the main file of a data type.
+
+    Parameters
+    ----------
+    datagrabber : DataGrabber-like object
+        The DataGrabber.
+    data_type : str
+        The data type.
+
+    Returns
+    -------
+    list of str or None
+        The nested data types (e.g., ``["confounds"]``) if the DataGrabber
+        only grabs them, without the main file of the data type, or None if
+        it grabs the main file.
+
+    """
+    patterns = getattr(datagrabber, "patterns", None)
+    if patterns is None:
+        return None
+    dtype_val = patterns.get(data_type)
+    # Data types with a list of patterns (e.g., Warp) have their main files
+    if dtype_val is None or isinstance(dtype_val, list):
+        return None
+    if "pattern" in dtype_val:
+        return None
+    return [k for k, v in dtype_val.items() if isinstance(v, dict)]
+
+
 @register_datagrabber
 class MultipleDataGrabber(BaseDataGrabber):
     """Concrete implementation for multi sourced data fetching.
 
     Implements a DataGrabber which can be used to fetch data from multiple
-    DataGrabbers.
+    DataGrabbers, e.g., from different datasets:
+
+    * Each data type is grabbed by one DataGrabber (set with their ``types``),
+      except for the DataGrabbers that only grab the nested data types of a
+      data type (without its main file, e.g., only the ``confounds`` of
+      ``BOLD``), which replace the nested data types of the other one.
+    * The element keys are the ones of the first DataGrabber, followed by the
+      other keys of the other DataGrabbers, which can have different keys
+      (e.g., ``subject`` and ``task`` for ``BOLD``, and only ``subject`` for
+      ``VBM_GM``). The keys with the same name must have the same values in
+      all the DataGrabbers.
+    * The elements are the ones available in all the DataGrabbers, joined on
+      their common keys, and each DataGrabber grabs an element with its keys.
 
     Parameters
     ----------
@@ -33,62 +83,66 @@ class MultipleDataGrabber(BaseDataGrabber):
     Raises
     ------
     RuntimeError
-        If ``datagrabbers`` have different element keys or
-        overlapping data types or nested data types.
+        If more than one DataGrabber grabs the main file of a data type or
+        the same nested data type without the main file.
 
     """
 
-    def __init__(self, datagrabbers: list[DataGrabberLike], **kwargs) -> None:
-        # Check datagrabbers consistency
-        # Check for same element keys
-        first_keys = datagrabbers[0].get_element_keys()
-        for dg in datagrabbers[1:]:
-            if dg.get_element_keys() != first_keys:
+    model_config = ConfigDict(extra="allow")
+
+    datagrabbers: list[
+        DataGrabberLike | PatternDataGrabber | PatternDataladDataGrabber
+    ]
+    types: Annotated[
+        DataType | list[DataType], BeforeValidator(ensure_list)
+    ] = []  # noqa: RUF012
+    datadir: Path = Path(".")
+
+    def validate_datagrabber_params(self) -> None:
+        """Run extra logical validation for datagrabber."""
+        # Each data type is grabbed by one DataGrabber, except for the nested
+        # data types grabbed without the main file, which can only be
+        # grabbed by one of these DataGrabbers
+        main = defaultdict(list)
+        nested = defaultdict(list)
+        for i, dg in enumerate(self.datagrabbers):
+            for data_type in dg.get_types():
+                nested_types = _nested_types_only(dg, data_type)
+                if nested_types is None:
+                    main[data_type].append(i)
+                else:
+                    for nested_type in nested_types:
+                        nested[(data_type, nested_type)].append(i)
+        for data_type, dgs in main.items():
+            if len(dgs) > 1:
                 raise_error(
-                    msg="DataGrabbers have different element keys",
+                    msg=(
+                        f"Data type `{data_type}` is grabbed by more than one "
+                        f"DataGrabber ({dgs}), it must be removed from the "
+                        "`types` of all but one of them"
+                    ),
                     klass=RuntimeError,
                 )
-        # Check for no overlapping types (and nested data types)
-        types = [x for dg in datagrabbers for x in dg.get_types()]
-        if len(types) != len(set(types)):
-            if all(hasattr(dg, "patterns") for dg in datagrabbers):
-                first_patterns = datagrabbers[0].patterns
-                for dg in datagrabbers[1:]:
-                    for data_type in set(types):
-                        dtype_pattern = dg.patterns.get(data_type)
-                        if dtype_pattern is None:
-                            continue
-                        # Check if first-level keys of data type are same
-                        if (
-                            dtype_pattern.keys()
-                            == first_patterns[data_type].keys()
-                        ):
-                            raise_error(
-                                msg=(
-                                    "DataGrabbers have overlapping mandatory "
-                                    "and / or optional key(s) for data type: "
-                                    f"`{data_type}`"
-                                ),
-                                klass=RuntimeError,
-                            )
-            else:
-                # Can't check further
+        for (data_type, nested_type), dgs in nested.items():
+            if len(dgs) > 1:
                 raise_error(
-                    msg="DataGrabbers have overlapping types",
+                    msg=(
+                        f"Nested data type `{data_type}.{nested_type}` is "
+                        f"grabbed by more than one DataGrabber ({dgs})"
+                    ),
                     klass=RuntimeError,
                 )
-        self._datagrabbers = datagrabbers
 
-    def __getitem__(self, element: Union[str, tuple]) -> dict:
+    def __getitem__(self, element: Element) -> dict:
         """Implement indexing.
 
         Parameters
         ----------
-        element : str or tuple
+        element : `Element`
             The element to be indexed. If one string is provided, it is
             assumed to be a tuple with only one item. If a tuple is provided,
-            each item in the tuple is the value for the replacement string
-            specified in "replacements".
+            each item in the tuple is the value for the element key (see
+            :meth:`.get_element_keys`).
 
         Returns
         -------
@@ -97,57 +151,104 @@ class MultipleDataGrabber(BaseDataGrabber):
             specified element.
 
         """
-
-        out = {}
+        if not isinstance(element, tuple):
+            element = (element,)
+        named_element = dict(
+            zip(self.get_element_keys(), element, strict=True)
+        )
+        # Grab the element with the keys of each DataGrabber
+        outs = []
         metas = []
-        for dg in self._datagrabbers:
-            t_out = dg[element]
-            deep_update(out, t_out)
+        for dg in self.datagrabbers:
+            t_element = tuple(named_element[k] for k in dg.get_element_keys())
+            outs.append((dg, dg[t_element]))
             # Now get the meta for this datagrabber
             t_meta = {}
             dg.update_meta(t_meta, "datagrabber")
             # Store all the sub-datagrabbers meta
             metas.append(t_meta["meta"]["datagrabber"])
 
+        # Add the data types with their main file first, and then the nested
+        # data types grabbed without the main file, which replace the ones
+        # grabbed with the main file
+        out = {}
+        for nested_only in (False, True):
+            for dg, t_out in outs:
+                for data_type, value in t_out.items():
+                    t_nested_only = (
+                        _nested_types_only(dg, data_type) is not None
+                    )
+                    if t_nested_only != nested_only:
+                        continue
+                    if data_type in out:
+                        deep_update(out[data_type], value)
+                    else:
+                        out[data_type] = value
+
         # Update all the metas again
         for kind in out:
             to_update = out[kind]
             if not isinstance(to_update, list):
                 to_update = [to_update]
+            element_keys = self.get_type_element_keys(kind)
             for t_kind in to_update:
                 self.update_meta(t_kind, "datagrabber")
                 t_kind["meta"]["datagrabber"]["datagrabbers"] = metas
+                t_kind["meta"]["element"] = named_element
+                t_kind["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.
+
+        Parameters
+        ----------
+        data_type : str
+            The data type.
+
+        Returns
+        -------
+        list of str
+            The element keys that the data of the data type depends on in all
+            the DataGrabbers grabbing it (e.g., the one grabbing the main file
+            and the one grabbing only its nested data types).
+
+        """
+        keys = set()
+        for dg in self.datagrabbers:
+            if data_type in dg.get_types():
+                keys.update(dg.get_type_element_keys(data_type))
+        return [k for k in self.get_element_keys() if k in keys]
+
     def __enter__(self) -> "MultipleDataGrabber":
         """Implement context entry."""
-        for dg in self._datagrabbers:
+        for dg in self.datagrabbers:
             dg.__enter__()
         return self
 
     def __exit__(self, exc_type, exc_value, exc_traceback) -> None:
         """Implement context exit."""
-        for dg in self._datagrabbers:
+        for dg in self.datagrabbers:
             dg.__exit__(exc_type, exc_value, exc_traceback)
 
-    # TODO: return type should be List[List[str]], but base type is List[str]
     def get_types(self) -> list[str]:
         """Get types.
 
         Returns
         -------
-        list of list of str
+        list of str
             The types of data to be grabbed.
 
         """
-        types = [x for dg in self._datagrabbers for x in dg.get_types()]
-        return types
+        types = [x for dg in self.datagrabbers for x in dg.get_types()]
+        return list(dict.fromkeys(types))
 
     def get_element_keys(self) -> list[str]:
         """Get element keys.
 
         For each item in the ``element`` tuple passed to ``__getitem__()``,
-        this method returns the corresponding key(s).
+        this method returns the corresponding key(s): the keys of the first
+        DataGrabber, followed by the other keys of the other DataGrabbers.
 
         Returns
         -------
@@ -155,11 +256,16 @@ class MultipleDataGrabber(BaseDataGrabber):
             The element keys.
 
         """
-        return self._datagrabbers[0].get_element_keys()
+        keys = [k for dg in self.datagrabbers for k in dg.get_element_keys()]
+        return list(dict.fromkeys(keys))
 
     def get_elements(self) -> list:
         """Get elements.
 
+        The elements of the DataGrabbers are joined on their common keys, as
+        the elements of the data types of a :class:`.PatternDataGrabber`, so
+        the elements are the ones available in all the DataGrabbers.
+
         Returns
         -------
         list
@@ -169,11 +275,25 @@ class MultipleDataGrabber(BaseDataGrabber):
             related DataGrabbers.
 
         """
-        all_elements = [dg.get_elements() for dg in self._datagrabbers]
-        elements = set(all_elements[0])
-        for s in all_elements[1:]:
-            elements.intersection_update(s)
-        return list(elements)
+        elements = None
+        keys = []
+        for dg in self.datagrabbers:
+            t_keys = dg.get_element_keys()
+            t_elements = {
+                e if isinstance(e, tuple) else (e,) for e in dg.get_elements()
+            }
+            if elements is None:
+                elements, keys = t_elements, t_keys
+            else:
+                elements, keys = _join_elements(
+                    elements, keys, t_elements, t_keys
+                )
+        if elements is None:
+            return []
+        out = sorted(elements)
+        if len(keys) == 1:
+            out = [element[0] for element in out]
+        return out
 
     def get_item(self, **_: dict) -> dict[str, dict]:
         """Get the specified item from the dataset.