Replace pybaum with optree and add paths to PythonNode names. (#396)

Author: tobiasraabe

.pre-commit-config.yaml

@@ -49,7 +49,7 @@ repos:
     rev: 1.5.0
     hooks:
     -   id: interrogate
-        args: [-v, --fail-under=75]
+        args: [-vv, --fail-under=75]
         exclude: ^(tests/|docs/|scripts/)
 -   repo: https://github.com/pre-commit/mirrors-mypy
     rev: 'v1.3.0'

docs/rtd_environment.yml

@@ -25,7 +25,7 @@ dependencies:
   - click-default-group
   - networkx >=2.4
   - pluggy
-  - pybaum >=0.1.1
+  - optree >=0.9
   - pexpect
   - rich
   - sqlalchemy >=1.4.36

docs/source/changes.md

@@ -15,6 +15,8 @@ releases are available on [PyPI](https://pypi.org/project/pytask) and
 - {pull}`394` allows to add products with {obj}`typing.Annotation` and
   {obj}`~pytask.Product`.
 - {pull}`395` refactors all occurrences of pybaum to {mod}`_pytask.tree_util`.
+- {pull}`396` replaces pybaum with optree and adds paths to the name of
+  {class}`pytask.PythonNode`'s allowing for better hashing.
 
 ## 0.3.2 - 2023-06-07
 

docs/source/reference_guides/api.md

@@ -247,6 +247,12 @@ Then, different kinds of nodes can be implemented.
 
 ```{eval-rst}
 .. autoclass:: pytask.FilePathNode
+    :members:
+```
+
+```{eval-rst}
+.. autoclass:: pytask.PythonNode
+    :members:
 ```
 
 To parse dependencies and products from nodes, use the following functions.

environment.yml

@@ -16,7 +16,7 @@ dependencies:
   - click-default-group
   - networkx >=2.4
   - pluggy
-  - pybaum >=0.1.1
+  - optree >=0.9
   - rich
   - sqlalchemy >=1.4.36
   - tomli >=1.0.0

setup.cfg

@@ -34,9 +34,9 @@ install_requires =
     click
     click-default-group
     networkx>=2.4
+    optree>=0.9
     packaging
     pluggy
-    pybaum>=0.1.1
     rich
     sqlalchemy>=1.4.36
     tomli>=1.0.0

src/_pytask/clean.py

@@ -29,7 +29,7 @@
 from _pytask.session import Session
 from _pytask.shared import to_list
 from _pytask.traceback import render_exc_info
-from _pytask.tree_util import tree_just_yield
+from _pytask.tree_util import tree_leaves
 from attrs import define
 
 
@@ -218,7 +218,7 @@ def _yield_paths_from_task(task: Task) -> Generator[Path, None, None]:
     """Yield all paths attached to a task."""
     yield task.path
     for attribute in ("depends_on", "produces"):
-        for node in tree_just_yield(getattr(task, attribute)):
+        for node in tree_leaves(getattr(task, attribute)):
             if hasattr(node, "path") and isinstance(node.path, Path):
                 yield node.path
 

src/_pytask/collect.py

@@ -11,9 +11,7 @@
 from typing import Generator
 from typing import Iterable
 
-from _pytask.collect_utils import depends_on
 from _pytask.collect_utils import parse_dependencies_from_task_function
-from _pytask.collect_utils import parse_nodes
 from _pytask.collect_utils import parse_products_from_task_function
 from _pytask.config import hookimpl
 from _pytask.config import IS_FILE_SYSTEM_CASE_SENSITIVE
@@ -22,7 +20,9 @@
 from _pytask.console import format_task_id
 from _pytask.exceptions import CollectionError
 from _pytask.mark_utils import has_mark
+from _pytask.models import NodeInfo
 from _pytask.nodes import FilePathNode
+from _pytask.nodes import MetaNode
 from _pytask.nodes import PythonNode
 from _pytask.nodes import Task
 from _pytask.outcomes import CollectionOutcome
@@ -169,13 +169,7 @@ def pytask_collect_task(
 
     """
     if (name.startswith("task_") or has_mark(obj, "task")) and callable(obj):
-        if has_mark(obj, "depends_on"):
-            nodes = parse_nodes(session, path, name, obj, depends_on)
-            dependencies = {"depends_on": nodes}
-        else:
-            dependencies = parse_dependencies_from_task_function(
-                session, path, name, obj
-            )
+        dependencies = parse_dependencies_from_task_function(session, path, name, obj)
 
         products = parse_products_from_task_function(session, path, name, obj)
 
@@ -210,9 +204,7 @@ def pytask_collect_task(
 
 
 @hookimpl(trylast=True)
-def pytask_collect_node(
-    session: Session, path: Path, node: str | Path
-) -> FilePathNode | PythonNode:
+def pytask_collect_node(session: Session, path: Path, node_info: NodeInfo) -> MetaNode:
     """Collect a node of a task as a :class:`pytask.nodes.FilePathNode`.
 
     Strings are assumed to be paths. This might be a strict assumption, but since this
@@ -222,17 +214,18 @@ def pytask_collect_node(
     ``trylast=True`` might be necessary if other plugins try to parse strings themselves
     like a plugin for downloading files which depends on URLs given as strings.
 
-    Parameters
-    ----------
-    session : _pytask.session.Session
-        The session.
-    path : Union[str, pathlib.Path]
-        The path to file where the task and node are specified.
-    node : Union[str, pathlib.Path]
-        The value of the node which can be a str, a path or anything which cannot be
-        handled by this function.
-
     """
+    node = node_info.value
+
+    if isinstance(node, PythonNode):
+        if not node.name:
+            suffix = "-" + "-".join(map(str, node_info.path)) if node_info.path else ""
+            node.name = node_info.arg_name + suffix
+        return node
+
+    if isinstance(node, MetaNode):
+        return node
+
     if isinstance(node, Path):
         if not node.is_absolute():
             node = path.parent.joinpath(node)
@@ -251,7 +244,10 @@ def pytask_collect_node(
                 raise ValueError(_TEMPLATE_ERROR.format(node, case_sensitive_path))
 
         return FilePathNode.from_path(node)
-    return PythonNode(value=node)
+
+    suffix = "-" + "-".join(map(str, node_info.path)) if node_info.path else ""
+    node_name = node_info.arg_name + suffix
+    return PythonNode(value=node, name=node_name)
 
 
 def _not_ignored_paths(
@@ -263,18 +259,6 @@ def _not_ignored_paths(
     directories, all subsequent files and folders are considered, but one level after
     another, so that files of ignored folders are not checked.
 
-    Parameters
-    ----------
-    paths : Iterable[pathlib.Path]
-        List of paths from which tasks are collected.
-    session : _pytask.session.Session
-        The session.
-
-    Yields
-    ------
-    path : pathlib.Path
-        A path which is not ignored.
-
     """
     for path in paths:
         if not session.hook.pytask_ignore_collect(path=path, config=session.config):
@@ -287,11 +271,7 @@ def _not_ignored_paths(
 
 @hookimpl(trylast=True)
 def pytask_collect_modify_tasks(tasks: list[Task]) -> None:
-    """Given all tasks, assign a short uniquely identifiable name to each task.
-
-    The shorter ids are necessary to display
-
-    """
+    """Given all tasks, assign a short uniquely identifiable name to each task."""
     id_to_short_id = _find_shortest_uniquely_identifiable_name_for_tasks(tasks)
     for task in tasks:
         short_id = id_to_short_id[task.name]

src/_pytask/collect_command.py

@@ -26,7 +26,7 @@
 from _pytask.path import relative_to
 from _pytask.pluginmanager import get_plugin_manager
 from _pytask.session import Session
-from _pytask.tree_util import tree_just_flatten
+from _pytask.tree_util import tree_leaves
 from rich.text import Text
 from rich.tree import Tree
 
@@ -126,10 +126,10 @@ def _find_common_ancestor_of_all_nodes(
         if show_nodes:
             all_paths.extend(
                 x.path
-                for x in tree_just_flatten(task.depends_on)
+                for x in tree_leaves(task.depends_on)
                 if isinstance(x, FilePathNode)
             )
-            all_paths.extend(x.path for x in tree_just_flatten(task.produces))
+            all_paths.extend(x.path for x in tree_leaves(task.produces))
 
     common_ancestor = find_common_ancestor(*all_paths, *paths)
 
@@ -150,7 +150,7 @@ def _organize_tasks(tasks: list[Task]) -> dict[Path, list[Task]]:
 
     sorted_dict = {}
     for k in sorted(dictionary):
-        sorted_dict[k] = sorted(dictionary[k], key=lambda x: x.path)
+        sorted_dict[k] = sorted(dictionary[k], key=lambda x: x.name)
 
     return sorted_dict
 
@@ -202,36 +202,24 @@ def _print_collected_tasks(
             )
 
             if show_nodes:
-                file_path_nodes = [
-                    i
-                    for i in tree_just_flatten(task.depends_on)
-                    if isinstance(i, FilePathNode)
-                ]
-                sorted_nodes = sorted(file_path_nodes, key=lambda x: x.path)
+                file_path_nodes = list(tree_leaves(task.depends_on))
+                sorted_nodes = sorted(file_path_nodes, key=lambda x: x.name)
                 for node in sorted_nodes:
-                    reduced_node_name = relative_to(node.path, common_ancestor)
-                    url_style = create_url_style_for_path(node.path, editor_url_scheme)
-                    task_branch.add(
-                        Text.assemble(
-                            FILE_ICON,
-                            "<Dependency ",
-                            Text(str(reduced_node_name), style=url_style),
-                            ">",
+                    if isinstance(node, FilePathNode):
+                        reduced_node_name = relative_to(node.path, common_ancestor)
+                        url_style = create_url_style_for_path(
+                            node.path, editor_url_scheme
                         )
-                    )
+                        text = Text(str(reduced_node_name), style=url_style)
+                    else:
+                        text = node.name
 
-                for node in sorted(
-                    tree_just_flatten(task.produces), key=lambda x: x.path
-                ):
+                    task_branch.add(Text.assemble(FILE_ICON, "<Dependency ", text, ">"))
+
+                for node in sorted(tree_leaves(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(
-                        Text.assemble(
-                            FILE_ICON,
-                            "<Product ",
-                            Text(str(reduced_node_name), style=url_style),
-                            ">",
-                        )
-                    )
+                    text = Text(str(reduced_node_name), style=url_style)
+                    task_branch.add(Text.assemble(FILE_ICON, "<Product ", text, ">"))
 
     console.print(tree)