Remove more type errors.

Author: tobiasraabe

src/_pytask/tree_util.py

@@ -4,9 +4,11 @@
 
 import functools
 from pathlib import Path
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import TypeVar
 
 import optree
-from optree import PyTree
 from optree import tree_flatten_with_path as _optree_tree_flatten_with_path
 from optree import tree_leaves as _optree_tree_leaves
 from optree import tree_map as _optree_tree_map
@@ -23,6 +25,19 @@
     "tree_structure",
 ]
 
+_T = TypeVar("_T")
+
+if TYPE_CHECKING:
+    # 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.
+    # See: https://github.com/metaopt/optree/issues/251
+    PyTree = (
+        _T | tuple["PyTree[_T]", ...] | list["PyTree[_T]"] | dict[Any, "PyTree[_T]"]
+    )
+else:
+    from optree import PyTree
+
 assert optree.__file__ is not None
 TREE_UTIL_LIB_DIRECTORY = Path(optree.__file__).parent
 

tests/test_collect_command.py

@@ -306,7 +306,7 @@ def task_example_2(path=Path("in_2.txt"), produces=Path("out_2.txt")): ...
 
 @define
 class Node:
-    path: Path
+    path: str
 
     def state(self): ...
 
@@ -321,8 +321,8 @@ def test_print_collected_tasks_without_nodes(capsys):
                 base_name="function",
                 path=Path("task_path.py"),
                 function=function,
-                depends_on={"depends_on": Node("in.txt")},  # type: ignore[arg-type]
-                produces={"produces": Node("out.txt")},  # type: ignore[arg-type]
+                depends_on={"depends_on": Node("in.txt")},
+                produces={"produces": Node("out.txt")},
             )
         ]
     }
@@ -343,8 +343,8 @@ def test_print_collected_tasks_with_nodes(capsys):
                 base_name="function",
                 path=Path("task_path.py"),
                 function=function,
-                depends_on={"depends_on": PathNode(name="in.txt", path=Path("in.txt"))},  # type: ignore[arg-type]
-                produces={"produces": PathNode(name="out.txt", path=Path("out.txt"))},  # type: ignore[arg-type]
+                depends_on={"depends_on": PathNode(name="in.txt", path=Path("in.txt"))},
+                produces={"produces": PathNode(name="out.txt", path=Path("out.txt"))},
             )
         ]
     }
@@ -366,10 +366,10 @@ def test_find_common_ancestor_of_all_nodes(show_nodes, expected_add):
             base_name="function",
             path=Path.cwd() / "src" / "task_path.py",
             function=function,
-            depends_on={  # type: ignore[arg-type]
+            depends_on={
                 "depends_on": PathNode.from_path(Path.cwd() / "src" / "in.txt")
             },
-            produces={  # type: ignore[arg-type]
+            produces={
                 "produces": PathNode.from_path(
                     Path.cwd().joinpath("..", "bld", "out.txt").resolve()
                 )