Add node-focused typing tests for tree utilities

pyproject.toml

@@ -62,6 +62,7 @@ docs = [
 docs-live = ["sphinx-autobuild>=2024.10.3"]
 plugin-list = ["httpx>=0.27.0", "tabulate[widechars]>=0.9.0"]
 test = [
+    { include-group = "typing" },
     "cloudpickle>=3.0.0",
     "deepdiff>=7.0.0",
     # nbmake requires pywin32 on Windows, which has no wheels for Python 3.14 yet
@@ -76,8 +77,11 @@ test = [
     "coiled>=1.42.0; python_version < '3.14'",
     "networkx>=2.4.0",
     "pygraphviz>=1.12;platform_system=='Linux'",
+    "pytest-typing>=26.3.0",
+]
+typing = [
+    "ty>=0.0.84",
 ]
-typing = ["ty>=0.0.8"]
 
 [project.urls]
 Changelog = "https://pytask-dev.readthedocs.io/en/stable/changes.html"
@@ -159,6 +163,7 @@ force-single-line = true
 convention = "numpy"
 
 [tool.pytest.ini_options]
+typing_checkers = ["ty"]
 testpaths = ["src", "tests"]
 filterwarnings = [
     "ignore:'@pytask.mark.*. is deprecated:FutureWarning",

src/_pytask/execute.py

@@ -285,7 +285,10 @@ def pytask_execute_task(session: Session, task: PTask) -> bool:
             )
             raise ValueError(msg)
 
-        nodes = tree_leaves(task.produces["return"])
+        # Recursive tree inference can retain containers in the inferred leaf type.
+        nodes = cast(
+            "list[PNode | PProvisionalNode]", tree_leaves(task.produces["return"])
+        )
         values = structure_return.flatten_up_to(out)
         for node, value in zip(nodes, values, strict=False):
             if not isinstance(node, PProvisionalNode):
@@ -301,9 +304,9 @@ def pytask_execute_task_teardown(session: Session, task: PTask) -> None:
         return
 
     collect_provisional_products(session, task)
-    missing_nodes: list[Any] = [
-        node for node in tree_leaves(task.produces) if not node.state()
-    ]
+    # Provisional products have been replaced with concrete nodes at this point.
+    nodes = cast("list[PNode]", tree_leaves(task.produces))
+    missing_nodes: list[Any] = [node for node in nodes if not node.state()]
     if missing_nodes:
         paths = session.config["paths"]
         files = [format_node_name(i, paths).plain for i in missing_nodes]

src/_pytask/tree_util.py

@@ -28,8 +28,11 @@
 _T = TypeVar("_T")
 _U = TypeVar("_U")
 _K = TypeVar("_K")
+_NodeT = TypeVar("_NodeT", bound="TaskNode")
 
 if TYPE_CHECKING:
+    from _pytask.node_protocols import TaskNode
+
     # Use our own recursive type alias for static type checking.
     # optree's PyTree uses __class_getitem__ to generate Union types at runtime,
     # but type checkers like ty cannot evaluate these dynamic types properly.
@@ -46,6 +49,48 @@
 _pytree = optree.pytree.reexport(namespace="pytask")
 
 
+# Custom leaf predicates can make containers into leaves, so the node overloads
+# apply only when the default traversal is used.
+@overload
+def tree_leaves(
+    tree: _NodeT,
+    /,
+    is_leaf: None = None,
+) -> list[_NodeT]: ...
+
+
+@overload
+def tree_leaves(
+    tree: list[_NodeT],
+    /,
+    is_leaf: None = None,
+) -> list[_NodeT]: ...
+
+
+@overload
+def tree_leaves(
+    tree: tuple[_NodeT, ...],
+    /,
+    is_leaf: None = None,
+) -> list[_NodeT]: ...
+
+
+@overload
+def tree_leaves(
+    tree: dict[_K, _NodeT],
+    /,
+    is_leaf: None = None,
+) -> list[_NodeT]: ...
+
+
+@overload
+def tree_leaves(
+    tree: PyTree[_T],
+    /,
+    is_leaf: Callable[[Any], bool] | None = None,
+) -> list[_T]: ...
+
+
 def tree_leaves(
     tree: PyTree[_T],
     /,