Use pybaum for more flexible dependencies and products.

docs/rtd_environment.yml

@@ -25,6 +25,7 @@ dependencies:
   - networkx
   - pluggy
   - pony >=0.7.15
+  - pybaum
   - pexpect
   - rich
   - typing-extensions

docs/source/changes.rst

@@ -9,6 +9,8 @@ all releases are available on `PyPI <https://pypi.org/project/pytask>`_ and
 0.2.0 - 2022-xx-xx
 ------------------
 
+- :pull:`211` allows for flexible dependencies and products which can be any pytree of
+  native Python objects as supported by pybaum.
 - :pull:`227` implements ``task.kwargs`` as a new way for a task to hold parametrized
   arguments. It also implements :class:`_pytask.models.CollectionMetadata` to carry
   parametrized arguments to the task class.

docs/source/tutorials/how_to_define_dependencies_products.rst

@@ -85,8 +85,30 @@ Multiple dependencies and products
 ----------------------------------
 
 Most tasks have multiple dependencies or products. The easiest way to attach multiple
-dependencies or products to a task is to pass a :class:`dict`, :class:`list` or another
-iterator to the marker containing the paths.
+dependencies or products to a task is to pass a :class:`dict` (highly recommended),
+:class:`list` or another iterator to the marker containing the paths.
+
+To assign labels to dependencies or products, pass a dictionary. For example,
+
+.. code-block:: python
+
+    @pytask.mark.produces({"first": BLD / "data_0.pkl", "second": BLD / "data_1.pkl"})
+    def task_create_random_data(produces):
+        ...
+
+Then, use
+
+.. code-block:: pycon
+
+    >>> produces["first"]
+    BLD / "data_0.pkl"
+
+    >>> produces["second"]
+    BLD / "data_1.pkl"
+
+inside the task function.
+
+You can also use lists and other iterables.
 
 .. code-block:: python
 
@@ -102,8 +124,9 @@ where keys are the positions in the list.
     >>> produces
     {0: BLD / "data_0.pkl", 1: BLD / "data_1.pkl"}
 
-Why dictionaries and not lists? First, dictionaries with positions as keys behave very
-similar to lists and conversion between both is easy.
+Why does pytask recommend dictionaries and even converts lists to dictionaries? First,
+dictionaries with positions as keys behave very similar to lists and conversion between
+both is easy.
 
 .. tip::
 
@@ -113,47 +136,100 @@ Secondly, dictionaries use keys instead of positions which is more verbose and
 descriptive and does not assume a fixed ordering. Both attributes are especially
 desirable in complex projects.
 
-To assign labels to dependencies or products, pass a dictionary. For example,
+
+Multiple decorators
+-------------------
+
+You can also attach multiple decorators to a function which will be merged into a single
+dictionary. This might help you to group certain dependencies and apply them to multiple
+tasks.
 
 .. code-block:: python
 
-    @pytask.mark.produces({"first": BLD / "data_0.pkl", "second": BLD / "data_1.pkl"})
-    def task_create_random_data(produces):
+    common_dependencies = ["text_1.txt", "text_2.txt"]
+
+
+    @pytask.mark.depends_on(common_dependencies)
+    @pytask.mark.depends_on("text_3.txt")
+    def task_example():
         ...
 
-Then, use
 
-.. code-block:: pycon
+Nested dependencies and products
+--------------------------------
 
-    >>> produces["first"]
-    BLD / "data_0.pkl"
+Dependencies and products are allowed to be nested containers consisting of tuples,
+lists, and dictionaries. In situations where you want more structure and are otherwise
+forced to flatten your inputs, this can be beneficial.
 
-    >>> produces["second"]
-    BLD / "data_1.pkl"
+Here is an example with a task which fits some model on data. It depends on a module
+containing the code for the model which is not actively used, but ensures that the task
+is rerun when the model is changed. And, it depends on data.
 
-inside the task function.
+.. code-block:: python
 
+    @pytask.mark.depends_on(
+        {
+            "model": [SRC / "models" / "model.py"],
+            "data": {"a": SRC / "data" / "a.pkl", "b": SRC / "data" / "b.pkl"},
+        }
+    )
+    @pytask.mark.produces(BLD / "models" / "fitted_model.pkl")
+    def task_fit_model():
+        ...
 
-Multiple decorators
--------------------
+It is also possible to merge nested containers. For example, you might want to reuse
+the dependency on models for other tasks as well.
 
-You can also attach multiple decorators to a function which will be merged into a single
-dictionary. This might help you to group certain dependencies and apply them to multiple
-tasks.
+.. code-block:: python
+
+    model_dependencies = pytask.mark.depends_on({"model": [SRC / "models" / "model.py"]})
+
+
+    @model_dependencies
+    @pytask.mark.depends_on(
+        {"data": {"a": SRC / "data" / "a.pkl", "b": SRC / "data" / "b.pkl"}}
+    )
+    @pytask.mark.produces(BLD / "models" / "fitted_model.pkl")
+    def task_fit_model():
+        ...
+
+In both cases, ``depends_on`` within the function will be
 
 .. code-block:: python
 
-    common_dependencies = ["text_1.txt", "text_2.txt"]
+    {
+        "model": [SRC / "models" / "model.py"],
+        "data": {"a": SRC / "data" / "a.pkl", "b": SRC / "data" / "b.pkl"},
+    }
 
+Tuples and lists are converted to dictionaries with integer keys. The innermost
+decorator is evaluated first.
 
-    @pytask.mark.depends_on(common_dependencies)
-    @pytask.mark.depends_on("text_3.txt")
-    def task_example():
+.. code-block:: python
+
+    @pytask.mark.depends_on([SRC / "models" / "model.py"])
+    @pytask.mark.depends_on([SRC / "data" / "a.pkl", SRC / "data" / "b.pkl"])
+    @pytask.mark.produces(BLD / "models" / "fitted_model.pkl")
+    def task_fit_model():
         ...
 
+would give
+
+.. code-block:: python
+
+    {0: SRC / "data" / "a.pkl", 1: SRC / "data" / "b.pkl", 2: SRC / "models" / "model.py"}
+
+.. seealso::
+
+    The general concept behind nested objects like tuples, lists, and dictionaries is
+    called pytrees and is more extensively explained in the `documentation of pybaum
+    <https://github.com/OpenSourceEconomics/pybaum>`_ which serves pytask under the
+    hood.
+
 
 References
 ----------
 
 .. [1] The official documentation for :mod:`pathlib`.
-.. [2] A guide for pathlib at `RealPython <https://realpython.com/python-pathlib/>`_.
+.. [2] A guide for pathlib by `realpython <https://realpython.com/python-pathlib/>`_.

environment.yml

@@ -22,6 +22,7 @@ dependencies:
   - networkx
   - pluggy
   - pony >=0.7.15
+  - pybaum
   - rich
   - typing-extensions
 

setup.cfg

@@ -41,6 +41,7 @@ install_requires =
     packaging
     pluggy
     pony>=0.7.15
+    pybaum
     rich
     typing-extensions
 python_requires = >=3.7

src/_pytask/clean.py

@@ -9,6 +9,7 @@
 from typing import Any
 from typing import Generator
 from typing import Iterable
+from typing import List
 from typing import TYPE_CHECKING
 
 import attr
@@ -26,6 +27,7 @@
 from _pytask.session import Session
 from _pytask.shared import get_first_non_none_value
 from _pytask.traceback import render_exc_info
+from pybaum.tree_util import tree_just_yield
 
 
 if TYPE_CHECKING:
@@ -190,7 +192,7 @@ def _yield_paths_from_task(task: MetaTask) -> Generator[Path, None, None]:
     """Yield all paths attached to a task."""
     yield task.path
     for attribute in ["depends_on", "produces"]:
-        for node in getattr(task, attribute).values():
+        for node in tree_just_yield(getattr(task, attribute)):
             if hasattr(node, "path") and isinstance(node.path, Path):
                 yield node.path
 
@@ -234,7 +236,7 @@ class _RecursivePathNode:
     """
 
     path = attr.ib(type=Path)
-    sub_nodes = attr.ib(type="list[_RecursivePathNode]")
+    sub_nodes = attr.ib(type=List["_RecursivePathNode"])
     is_dir = attr.ib(type=bool)
     is_file = attr.ib(type=bool)
     is_unknown = attr.ib(type=bool)

src/_pytask/collect_command.py

@@ -25,6 +25,7 @@
 from _pytask.path import relative_to
 from _pytask.pluginmanager import get_plugin_manager
 from _pytask.session import Session
+from pybaum.tree_util import tree_just_flatten
 from rich.text import Text
 from rich.tree import Tree
 
@@ -125,13 +126,8 @@ def _find_common_ancestor_of_all_nodes(
     for task in tasks:
         all_paths.append(task.path)
         if show_nodes:
-            all_paths.extend(
-                [
-                    node.path
-                    for attr in ("depends_on", "produces")
-                    for node in getattr(task, attr).values()
-                ]
-            )
+            all_paths.extend(map(lambda x: x.path, tree_just_flatten(task.depends_on)))
+            all_paths.extend(map(lambda x: x.path, tree_just_flatten(task.produces)))
 
     common_ancestor = find_common_ancestor(*all_paths, *paths)
 
@@ -201,7 +197,9 @@ def _print_collected_tasks(
             )
 
             if show_nodes:
-                for node in sorted(task.depends_on.values(), key=lambda x: x.path):
+                for node in sorted(
+                    tree_just_flatten(task.depends_on), key=lambda x: x.path
+                ):
                     reduced_node_name = relative_to(node.path, common_ancestor)
                     url_style = create_url_style_for_path(node.path, editor_url_scheme)
                     task_branch.add(
@@ -213,7 +211,9 @@ def _print_collected_tasks(
                         )
                     )
 
-                for node in sorted(task.produces.values(), key=lambda x: x.path):
+                for node in sorted(
+                    tree_just_flatten(task.produces), key=lambda x: x.path
+                ):
                     reduced_node_name = relative_to(node.path, common_ancestor)
                     url_style = create_url_style_for_path(node.path, editor_url_scheme)
                     task_branch.add(

src/_pytask/execute.py

@@ -29,6 +29,7 @@
 from _pytask.traceback import format_exception_without_traceback
 from _pytask.traceback import remove_traceback_from_exc_info
 from _pytask.traceback import render_exc_info
+from pybaum.tree_util import tree_map
 from rich.text import Text
 
 
@@ -157,13 +158,7 @@ def pytask_execute_task(task: MetaTask) -> bool:
     for arg_name in ("depends_on", "produces"):
         if arg_name in func_arg_names:
             attribute = getattr(task, arg_name)
-            kwargs[arg_name] = (
-                attribute[0].value
-                if len(attribute) == 1
-                and 0 in attribute
-                and not task.keep_dict[arg_name]
-                else {name: node.value for name, node in attribute.items()}
-            )
+            kwargs[arg_name] = tree_map(lambda x: x.value, attribute)
 
     task.execute(**kwargs)
     return True