Handle local UPath protocols consistently

Author: tobiasraabe

src/_pytask/collect.py

@@ -43,6 +43,7 @@
 from _pytask.path import find_case_sensitive_path
 from _pytask.path import import_path
 from _pytask.path import is_non_local_path
+from _pytask.path import normalize_local_upath
 from _pytask.path import shorten_path
 from _pytask.pluginmanager import hookimpl
 from _pytask.reports import CollectionReport
@@ -456,6 +457,9 @@ def pytask_collect_node(  # noqa: C901, PLR0912
             node.name = create_name_of_python_node(node_info)
         return node
 
+    if isinstance(node, PPathNode):
+        node.path = normalize_local_upath(node.path)
+
     if (
         isinstance(node, PPathNode)
         and not is_non_local_path(node.path)
@@ -492,6 +496,9 @@ def pytask_collect_node(  # noqa: C901, PLR0912
             node.name = create_name_of_python_node(node_info)
         return node
 
+    if isinstance(node, UPath):  # pragma: no cover
+        node = normalize_local_upath(node)
+
     if isinstance(node, UPath):  # pragma: no cover
         if not node.protocol:
             node = Path(node)

src/_pytask/collect_command.py

@@ -31,6 +31,7 @@
 from _pytask.outcomes import ExitCode
 from _pytask.path import find_common_ancestor
 from _pytask.path import is_non_local_path
+from _pytask.path import normalize_local_upath
 from _pytask.path import relative_to
 from _pytask.pluginmanager import hookimpl
 from _pytask.pluginmanager import storage
@@ -126,12 +127,12 @@ def _find_common_ancestor_of_all_nodes(
         all_paths.append(task.path)
         if show_nodes:
             all_paths.extend(
-                x.path
+                normalize_local_upath(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
+                normalize_local_upath(x.path)
                 for x in tree_leaves(task.produces)
                 if isinstance(x, PPathNode) and not is_non_local_path(x.path)
             )

src/_pytask/path.py

@@ -28,11 +28,15 @@
     "hash_path",
     "import_path",
     "is_non_local_path",
+    "normalize_local_upath",
     "relative_to",
     "shorten_path",
 ]
 
 
+_LOCAL_UPATH_PROTOCOLS = frozenset(("", "file", "local"))
+
+
 def relative_to(path: Path, source: Path, *, include_source: bool = True) -> Path:
     """Make a path relative to another path.
 
@@ -61,7 +65,14 @@ def relative_to(path: Path, source: Path, *, include_source: bool = True) -> Pat
 
 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)
+    return isinstance(path, UPath) and path.protocol not in _LOCAL_UPATH_PROTOCOLS
+
+
+def normalize_local_upath(path: Path) -> Path:
+    """Convert local `UPath` variants to a stdlib `Path`."""
+    if isinstance(path, UPath) and path.protocol in {"file", "local"}:
+        return Path(path.path)
+    return path
 
 
 def find_closest_ancestor(
@@ -443,6 +454,9 @@ def shorten_path(path: Path, paths: Sequence[Path]) -> str:
     if is_non_local_path(path):
         return path.as_posix()
 
+    path = normalize_local_upath(path)
+    paths = [normalize_local_upath(p) for p in paths]
+
     ancestor = find_closest_ancestor(path, paths)
     if ancestor is None:
         try:

tests/test_collect.py

@@ -214,6 +214,31 @@ def test_pytask_collect_remote_path_node_keeps_uri_name():
     assert result.name == "s3://bucket/file.pkl"
 
 
+@pytest.mark.parametrize("protocol", ["file", "local"])
+def test_pytask_collect_local_upath_protocol_node_is_shortened(tmp_path, protocol):
+    upath = pytest.importorskip("upath")
+
+    session = Session.from_config(
+        {"check_casing_of_paths": False, "paths": (tmp_path,), "root": tmp_path}
+    )
+
+    result = pytask_collect_node(
+        session,
+        tmp_path,
+        NodeInfo(
+            arg_name="path",
+            path=(),
+            value=PickleNode(path=upath.UPath(f"{protocol}://{tmp_path}/file.pkl")),
+            task_path=tmp_path / "task_example.py",
+            task_name="task_example",
+        ),
+    )
+
+    assert isinstance(result, PPathNode)
+    assert result.path == tmp_path / "file.pkl"
+    assert result.name == f"{tmp_path.name}/file.pkl"
+
+
 @pytest.mark.skipif(
     sys.platform != "win32", reason="Only works on case-insensitive file systems."
 )

tests/test_collect_command.py

@@ -421,6 +421,32 @@ def task_example(
     assert "s3://bucket/in.pkl" in result.output
 
 
+@pytest.mark.parametrize("protocol", ["file", "local"])
+def test_collect_task_with_local_upath_protocol_node(runner, tmp_path, protocol):
+    pytest.importorskip("upath")
+
+    source = f"""
+    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("{protocol}://{tmp_path.as_posix()}/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 f"{tmp_path.name}/in.pkl" in result.output
+
+
 def test_python_node_is_collected(runner, tmp_path):
     source = """
     from pytask import Product

tests/test_path.py

@@ -19,6 +19,7 @@
 from _pytask.path import find_case_sensitive_path
 from _pytask.path import find_closest_ancestor
 from _pytask.path import find_common_ancestor
+from _pytask.path import is_non_local_path
 from _pytask.path import relative_to
 from _pytask.path import shorten_path
 from pytask.path import import_path
@@ -119,6 +120,16 @@ def test_shorten_path_keeps_non_local_uri():
     assert shorten_path(path, [Path.cwd()]) == "s3://bucket/file.pkl"
 
 
+@pytest.mark.parametrize("protocol", ["file", "local"])
+def test_shorten_path_treats_local_upath_protocols_as_local(tmp_path, protocol):
+    upath = pytest.importorskip("upath")
+
+    path = upath.UPath(f"{protocol}://{tmp_path.as_posix()}/file.pkl")
+
+    assert not is_non_local_path(path)
+    assert shorten_path(path, [tmp_path]) == f"{tmp_path.name}/file.pkl"
+
+
 @pytest.mark.skipif(sys.platform != "win32", reason="Only works on Windows.")
 @pytest.mark.parametrize(
     ("path", "existing_paths", "expected"),