Fix remote UPath collection and node display

Author: tobiasraabe

src/_pytask/collect_command.py

@@ -30,6 +30,7 @@
 from _pytask.node_protocols import PTaskWithPath
 from _pytask.outcomes import ExitCode
 from _pytask.path import find_common_ancestor
+from _pytask.path import is_non_local_path
 from _pytask.path import relative_to
 from _pytask.pluginmanager import hookimpl
 from _pytask.pluginmanager import storage
@@ -125,10 +126,14 @@ def _find_common_ancestor_of_all_nodes(
         all_paths.append(task.path)
         if show_nodes:
             all_paths.extend(
-                x.path for x in tree_leaves(task.depends_on) if isinstance(x, PPathNode)
+                x.path
+                for x in tree_leaves(task.depends_on)
+                if isinstance(x, PPathNode) and not is_non_local_path(x.path)
             )
             all_paths.extend(
-                x.path for x in tree_leaves(task.produces) if isinstance(x, PPathNode)
+                x.path
+                for x in tree_leaves(task.produces)
+                if isinstance(x, PPathNode) and not is_non_local_path(x.path)
             )
 
     return find_common_ancestor(*all_paths, *paths)

src/_pytask/path.py

@@ -13,6 +13,8 @@
 from types import ModuleType
 from typing import TYPE_CHECKING
 
+from upath import UPath
+
 from _pytask._hashlib import file_digest
 from _pytask.cache import Cache
 
@@ -25,6 +27,7 @@
     "find_common_ancestor",
     "hash_path",
     "import_path",
+    "is_non_local_path",
     "relative_to",
     "shorten_path",
 ]
@@ -56,6 +59,11 @@ def relative_to(path: Path, source: Path, *, include_source: bool = True) -> Pat
     return Path(source_name, path.relative_to(source))
 
 
+def is_non_local_path(path: Path) -> bool:
+    """Return whether a path points to a non-local `UPath` resource."""
+    return isinstance(path, UPath) and bool(path.protocol)
+
+
 def find_closest_ancestor(
     path: Path, potential_ancestors: Sequence[Path]
 ) -> Path | None:
@@ -432,6 +440,9 @@ def shorten_path(path: Path, paths: Sequence[Path]) -> str:
     path from one path in ``session.config["paths"]`` to the node.
 
     """
+    if is_non_local_path(path):
+        return path.as_posix()
+
     ancestor = find_closest_ancestor(path, paths)
     if ancestor is None:
         try:

tests/test_collect.py

@@ -14,6 +14,7 @@
 from pytask import CollectionOutcome
 from pytask import ExitCode
 from pytask import NodeInfo
+from pytask import PickleNode
 from pytask import Session
 from pytask import Task
 from pytask import build
@@ -190,6 +191,29 @@ def test_pytask_collect_node(session, path, node_info, expected):
     assert str(result.load()) == str(expected)
 
 
+def test_pytask_collect_remote_path_node_keeps_uri_name():
+    upath = pytest.importorskip("upath")
+
+    session = Session.from_config(
+        {"check_casing_of_paths": False, "paths": (Path.cwd(),), "root": Path.cwd()}
+    )
+
+    result = pytask_collect_node(
+        session,
+        Path.cwd(),
+        NodeInfo(
+            arg_name="path",
+            path=(),
+            value=PickleNode(path=upath.UPath("s3://bucket/file.pkl")),
+            task_path=Path.cwd() / "task_example.py",
+            task_name="task_example",
+        ),
+    )
+
+    assert isinstance(result, PPathNode)
+    assert result.name == "s3://bucket/file.pkl"
+
+
 @pytest.mark.skipif(
     sys.platform != "win32", reason="Only works on case-insensitive file systems."
 )

tests/test_collect_command.py

@@ -396,6 +396,31 @@ def test_task_name_is_shortened(runner, tmp_path):
     assert "a/b/task_example.py::task_example" not in result.output
 
 
+def test_collect_task_with_remote_upath_node(runner, tmp_path):
+    pytest.importorskip("upath")
+
+    source = """
+    from pathlib import Path
+    from typing import Annotated
+
+    from upath import UPath
+
+    from pytask import PickleNode
+    from pytask import Product
+
+    def task_example(
+        data=PickleNode(path=UPath("s3://bucket/in.pkl")),
+        path: Annotated[Path, Product] = Path("out.txt"),
+    ): ...
+    """
+    tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source))
+
+    result = runner.invoke(cli, ["collect", "--nodes", tmp_path.as_posix()])
+
+    assert result.exit_code == ExitCode.OK
+    assert "s3://bucket/in.pkl" in result.output
+
+
 def test_python_node_is_collected(runner, tmp_path):
     source = """
     from pytask import Product

tests/test_path.py

@@ -20,6 +20,7 @@
 from _pytask.path import find_closest_ancestor
 from _pytask.path import find_common_ancestor
 from _pytask.path import relative_to
+from _pytask.path import shorten_path
 from pytask.path import import_path
 
 if TYPE_CHECKING:
@@ -110,6 +111,14 @@ def test_find_common_ancestor(path_1, path_2, expectation, expected):
         assert result == expected
 
 
+def test_shorten_path_keeps_non_local_uri():
+    upath = pytest.importorskip("upath")
+
+    path = upath.UPath("s3://bucket/file.pkl")
+
+    assert shorten_path(path, [Path.cwd()]) == "s3://bucket/file.pkl"
+
+
 @pytest.mark.skipif(sys.platform != "win32", reason="Only works on Windows.")
 @pytest.mark.parametrize(
     ("path", "existing_paths", "expected"),