Refactor pybaum to _pytask.tree_util. (#395)

Author: tobiasraabe

docs/source/changes.md

@@ -12,6 +12,9 @@ releases are available on [PyPI](https://pypi.org/project/pytask) and
   is not present.
 - {pull}`387` replaces pony with sqlalchemy.
 - {pull}`391` removes `@pytask.mark.parametrize`.
+- {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`.
 
 ## 0.3.2 - 2023-06-07
 

src/_pytask/clean.py

@@ -29,8 +29,8 @@
 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 attrs import define
-from pybaum.tree_util import tree_just_yield
 
 
 if TYPE_CHECKING:

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 pybaum.tree_util import tree_just_flatten
+from _pytask.tree_util import tree_just_flatten
 from rich.text import Text
 from rich.tree import Tree
 

src/_pytask/collect_utils.py

@@ -19,9 +19,9 @@
 from _pytask.nodes import PythonNode
 from _pytask.shared import find_duplicates
 from _pytask.task_utils import parse_keyword_arguments_from_signature_defaults
+from _pytask.tree_util import tree_map
 from attrs import define
 from attrs import field
-from pybaum.tree_util import tree_map
 from typing_extensions import Annotated
 from typing_extensions import get_origin
 

src/_pytask/dag.py

@@ -31,7 +31,7 @@
 from _pytask.shared import reduce_names_of_multiple_nodes
 from _pytask.shared import reduce_node_name
 from _pytask.traceback import render_exc_info
-from pybaum import tree_map
+from _pytask.tree_util import tree_map
 from rich.text import Text
 from rich.tree import Tree
 

src/_pytask/execute.py

@@ -33,7 +33,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 _pytask.tree_util import tree_map
 from rich.text import Text
 
 

src/_pytask/traceback.py

@@ -10,7 +10,7 @@
 
 import _pytask
 import pluggy
-import pybaum
+from _pytask.tree_util import TREE_UTIL_LIB_DIRECTORY
 from rich.traceback import Traceback
 
 
@@ -23,7 +23,6 @@
 
 
 _PLUGGY_DIRECTORY = Path(pluggy.__file__).parent
-_PYBAUM_DIRECTORY = Path(pybaum.__file__).parent
 _PYTASK_DIRECTORY = Path(_pytask.__file__).parent
 
 
@@ -93,7 +92,7 @@ def _is_internal_or_hidden_traceback_frame(
     path = Path(frame.tb_frame.f_code.co_filename)
     return any(
         root in path.parents
-        for root in (_PLUGGY_DIRECTORY, _PYBAUM_DIRECTORY, _PYTASK_DIRECTORY)
+        for root in (_PLUGGY_DIRECTORY, TREE_UTIL_LIB_DIRECTORY, _PYTASK_DIRECTORY)
     )
 
 

src/_pytask/tree_util.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import pybaum
+from pybaum.tree_util import tree_just_flatten
+from pybaum.tree_util import tree_just_yield
+from pybaum.tree_util import tree_map
+
+
+__all__ = [
+    "tree_just_flatten",
+    "tree_map",
+    "tree_just_yield",
+    "TREE_UTIL_LIB_DIRECTORY",
+]
+
+TREE_UTIL_LIB_DIRECTORY = Path(pybaum.__file__).parent

tests/test_pybaum.py tests/test_tree_util.pytests/test_pybaum.py renamed to tests/test_tree_util.py

@@ -1,11 +1,11 @@
-"""This module contains tests for pybaum and flexible dependencies and products."""
+"""This module contains tests for tree_util and flexible dependencies and products."""
 from __future__ import annotations
 
 import textwrap
 
 import pytest
 from _pytask.outcomes import ExitCode
-from pybaum import tree_map
+from _pytask.tree_util import tree_map
 from pytask import cli
 from pytask import main
 
@@ -55,11 +55,11 @@ def task_example():
 
 
 @pytest.mark.end_to_end()
-def test_profile_with_pybaum(tmp_path, runner):
+def test_profile_with_pytree(tmp_path, runner):
     source = """
     import time
     import pytask
-    from pybaum.tree_util import tree_just_flatten
+    from _pytask.tree_util import tree_just_flatten
 
     @pytask.mark.produces([{"out_1": "out_1.txt"}, {"out_2": "out_2.txt"}])
     def task_example(produces):