Support chaining for zip and tar paths (#440)

* upath._chain: fix unnesting for correct roundtripping

* tests: add chained tarfile tests

* upath: correct use of vfspath and .path and .__str__

* upath.implementations.cloud: fix .path override

* tests: add chained zip tests

* upath._chain: populate target_options if target_protocol is set

* upath._flavour: fix windows join behavior

Author: ap--

upath/_chain.py

@@ -3,6 +3,7 @@
 import sys
 import warnings
 from collections import defaultdict
+from collections import deque
 from collections.abc import MutableMapping
 from collections.abc import Sequence
 from collections.abc import Set
@@ -105,7 +106,26 @@ def replace(
         return type(self)(*segments, index=index)
 
     def to_list(self) -> list[ChainSegment]:
-        return list(self._segments)
+        """return a list of chain segments unnesting target_* segments"""
+        queue = deque(self._segments)
+        segments = []
+        while queue:
+            segment = queue.popleft()
+            if (
+                "target_protocol" in segment.storage_options
+                and "fo" in segment.storage_options
+            ):
+                storage_options = segment.storage_options.copy()
+                target_options = storage_options.pop("target_options", {})
+                target_protocol = storage_options.pop("target_protocol")
+                fo = storage_options.pop("fo")
+                queue.appendleft(ChainSegment(fo, target_protocol, target_options))
+                segments.append(
+                    ChainSegment(segment.path, segment.protocol, storage_options)
+                )
+            elif not segments or segment != segments[-1]:
+                segments.append(segment)
+        return segments
 
     @classmethod
     def from_list(cls, segments: list[ChainSegment], index: int = 0) -> Self:
@@ -124,9 +144,9 @@ def nest(self) -> ChainSegment:
                 urls = _prev
             _prev = urls
             if i == len(chain) - 1:
-                inkwargs = dict(**kw, **inkwargs)
+                inkwargs = {**kw, **inkwargs}
                 continue
-            inkwargs["target_options"] = dict(**kw, **inkwargs)
+            inkwargs["target_options"] = {**kw, **inkwargs}
             inkwargs["target_protocol"] = protocol
             inkwargs["fo"] = urls  # codespell:ignore fo
         urlpath, protocol, _ = chain[0]
@@ -186,6 +206,8 @@ def unchain(self, path: str, kwargs: dict[str, Any]) -> list[ChainSegment]:
                 kws.update(kwargs)
             kw = dict(**extra_kwargs)
             kw.update(kws)
+            if "target_protocol" in kw:
+                kw.setdefault("target_options", {})
             bit = flavour.strip_protocol(bit) or flavour.root_marker
             if (
                 protocol in {"blockcache", "filecache", "simplecache"}
@@ -229,7 +251,8 @@ def chain(self, segments: Sequence[ChainSegment]) -> tuple[str, dict[str, Any]]:
                 continue
             urlpaths.append(urlpath)
             # TODO: ensure roundtrip with unchain behavior
-            kwargs[segment.protocol] = segment.storage_options
+            if segment.storage_options:
+                kwargs[segment.protocol] = segment.storage_options
         return self.link.join(urlpaths), kwargs
 
 

upath/_flavour.py

@@ -287,10 +287,11 @@ def join(self, path: JoinablePathLike, *paths: JoinablePathLike) -> str:
         if not paths:
             return self.strip_protocol(path) or self.root_marker
         if self.local_file:
-            return os.path.join(
+            p = os.path.join(
                 self.strip_protocol(path),
                 *map(self.stringify_path, paths),
             )
+            return p if os.name != "nt" else p.replace("\\", "/")
         if self.netloc_is_anchor:
             drv, p0 = self.splitdrive(path)
             pN = list(map(self.stringify_path, paths))

upath/core.py

@@ -402,27 +402,38 @@ def __init__(
                     **args0._chain.nest().storage_options,
                     **storage_options,
                 }
-                str_args0 = str(args0)
+                str_args0 = args0.__vfspath__()
 
             else:
                 if hasattr(args0, "__fspath__") and args0.__fspath__ is not None:
                     str_args0 = args0.__fspath__()
+                elif hasattr(args0, "__vfspath__") and args0.__vfspath__ is not None:
+                    str_args0 = args0.__vfspath__()
+                elif isinstance(args0, str):
+                    str_args0 = args0
                 else:
-                    str_args0 = str(args0)
+                    raise TypeError(
+                        "argument should be a UPath, str, "
+                        f"or support __vfspath__ or __fspath__, not {type(args0)!r}"
+                    )
                 storage_options = type(self)._parse_storage_options(
                     str_args0, protocol, storage_options
                 )
-            if len(args) > 1:
-                str_args0 = WrappedFileSystemFlavour.from_protocol(protocol).join(
-                    str_args0, *args[1:]
-                )
         else:
             str_args0 = "."
 
         segments = chain_parser.unchain(
             str_args0, {"protocol": protocol, **storage_options}
         )
-        self._chain = Chain.from_list(segments)
+        chain = Chain.from_list(segments)
+        if len(args) > 1:
+            chain = chain.replace(
+                path=WrappedFileSystemFlavour.from_protocol(protocol).join(
+                    chain.active_path,
+                    *args[1:],
+                )
+            )
+        self._chain = chain
         self._chain_parser = chain_parser
         self._raw_urlpaths = args
         self._relative_base = None
@@ -638,9 +649,6 @@ def with_segments(self, *pathsegments: JoinablePathLike) -> Self:
         return new_instance
 
     def __str__(self) -> str:
-        return self.__vfspath__()
-
-    def __vfspath__(self) -> str:
         if self._relative_base is not None:
             active_path = self._chain.active_path
             stripped_base = self.parser.strip_protocol(
@@ -658,10 +666,19 @@ def __vfspath__(self) -> str:
         else:
             return self._chain_parser.chain(self._chain.to_list())[0]
 
+    def __vfspath__(self) -> str:
+        if self._relative_base is not None:
+            return self.__str__()
+        else:
+            return self.path
+
     def __repr__(self) -> str:
+        cls_name = type(self).__name__
+        path = self.__vfspath__()
         if self._relative_base is not None:
-            return f"<relative {type(self).__name__} {str(self)!r}>"
-        return f"{type(self).__name__}({self.path!r}, protocol={self._protocol!r})"
+            return f"<relative {cls_name} {path!r}>"
+        else:
+            return f"{cls_name}({path!r}, protocol={self._protocol!r})"
 
     # === JoinablePath overrides ======================================
 
@@ -704,9 +721,9 @@ def with_name(self, name: str) -> Self:
         split = self.parser.split
         if self.parser.sep in name:  # `split(name)[0]`
             raise ValueError(f"Invalid name {name!r}")
-        path = str(self)
-        path = path.removesuffix(split(path)[1]) + name
-        return self.with_segments(path)
+        _path = self.__vfspath__()
+        _path = _path.removesuffix(split(_path)[1]) + name
+        return self.with_segments(_path)
 
     @property
     def anchor(self) -> str:
@@ -780,7 +797,7 @@ def iterdir(self) -> Iterator[Self]:
                 continue
             # only want the path name with iterdir
             _, _, name = name.removesuffix(sep).rpartition(self.parser.sep)
-            yield base.with_segments(str(base), name)
+            yield base.with_segments(base.path, name)
 
     def __open_reader__(self) -> BinaryIO:
         return self.fs.open(self.path, mode="rb")
@@ -1045,7 +1062,7 @@ def glob(
             self = self.absolute()
         path_pattern = self.joinpath(pattern).path
         sep = self.parser.sep
-        base = self.fs._strip_protocol(self.path)
+        base = self.path
         for name in self.fs.glob(path_pattern):
             name = name.removeprefix(base).removeprefix(sep)
             yield self.joinpath(name)
@@ -1075,7 +1092,7 @@ def rglob(
         if _FSSPEC_HAS_WORKING_GLOB:
             r_path_pattern = self.joinpath("**", pattern).path
             sep = self.parser.sep
-            base = self.fs._strip_protocol(self.path)
+            base = self.path
             for name in self.fs.glob(r_path_pattern):
                 name = name.removeprefix(base).removeprefix(sep)
                 yield self.joinpath(name)
@@ -1084,7 +1101,7 @@ def rglob(
             path_pattern = self.joinpath(pattern).path
             r_path_pattern = self.joinpath("**", pattern).path
             sep = self.parser.sep
-            base = self.fs._strip_protocol(self.path)
+            base = self.path
             seen = set()
             for p in (path_pattern, r_path_pattern):
                 for name in self.fs.glob(p):
@@ -1134,7 +1151,7 @@ def __eq__(self, other: object) -> bool:
                 return False
 
         return (
-            self.path == other.path
+            self.__vfspath__() == other.__vfspath__()
             and self.protocol == other.protocol
             and self.storage_options == other.storage_options
         )
@@ -1145,29 +1162,27 @@ def __hash__(self) -> int:
         Note: in the future, if hash collisions become an issue, we
           can add `fsspec.utils.tokenize(storage_options)`
         """
-        if self._relative_base is not None:
-            return hash((self.protocol, str(self)))
-        return hash((self.protocol, self.path))
+        return hash((self.protocol, self.__vfspath__()))
 
     def __lt__(self, other: object) -> bool:
         if not isinstance(other, UPath) or self.parser is not other.parser:
             return NotImplemented
-        return self.path < other.path
+        return self.__vfspath__() < other.__vfspath__()
 
     def __le__(self, other: object) -> bool:
         if not isinstance(other, UPath) or self.parser is not other.parser:
             return NotImplemented
-        return self.path <= other.path
+        return self.__vfspath__() <= other.__vfspath__()
 
     def __gt__(self, other: object) -> bool:
         if not isinstance(other, UPath) or self.parser is not other.parser:
             return NotImplemented
-        return self.path > other.path
+        return self.__vfspath__() > other.__vfspath__()
 
     def __ge__(self, other: object) -> bool:
         if not isinstance(other, UPath) or self.parser is not other.parser:
             return NotImplemented
-        return self.path >= other.path
+        return self.__vfspath__() >= other.__vfspath__()
 
     def resolve(self, strict: bool = False) -> Self:
         if self._relative_base is not None:

upath/implementations/cloud.py

@@ -60,8 +60,8 @@ def root(self) -> str:
             return ""
         return self.parser.sep
 
-    def __vfspath__(self) -> str:
-        path = super().__vfspath__()
+    def __str__(self) -> str:
+        path = super().__str__()
         if self._relative_base is None:
             drive = self.parser.splitdrive(path)[0]
             if drive and path == f"{self.protocol}://{drive}":
@@ -71,7 +71,11 @@ def __vfspath__(self) -> str:
     @property
     def path(self) -> str:
         self_path = super().path
-        if self._relative_base is None and self.parser.sep not in self_path:
+        if (
+            self._relative_base is None
+            and self_path
+            and self.parser.sep not in self_path
+        ):
             return self_path + self.root
         return self_path
 

upath/tests/cases.py

@@ -512,7 +512,7 @@ def test_as_uri(self):
         # test that we can reconstruct the path from the uri
         p0 = self.path
         uri = p0.as_uri()
-        p1 = UPath(uri, **p0.fs.storage_options)
+        p1 = UPath(uri, **p0.storage_options)
         assert p0 == p1
 
     def test_protocol(self):

upath/tests/implementations/test_tar.py

@@ -82,3 +82,19 @@ def test_move_memory(self, clear_fsspec_memory_cache):
     @pytest.mark.skip(reason="Only testing read on TarPath")
     def test_move_into_memory(self, clear_fsspec_memory_cache):
         pass
+
+
+@pytest.fixture(scope="function")
+def tarred_testdir_file_in_memory(tarred_testdir_file, clear_fsspec_memory_cache):
+    p = UPath(tarred_testdir_file, protocol="file")
+    t = p.move(UPath("memory:///mytarfile.tar"))
+    assert t.protocol == "memory"
+    assert t.exists()
+    yield t.as_uri()
+
+
+class TestChainedTarPath(TestTarPath):
+
+    @pytest.fixture(autouse=True)
+    def path(self, tarred_testdir_file_in_memory):
+        self.path = UPath("tar://::memory:///mytarfile.tar")

upath/tests/implementations/test_zip.py

@@ -145,3 +145,25 @@ def test_write_text(self):
     @pytest.mark.skip(reason="fsspec zipfile filesystem is either read xor write mode")
     def test_fsspec_compat(self):
         pass
+
+
+@pytest.fixture(scope="function")
+def zipped_testdir_file_in_memory(zipped_testdir_file, clear_fsspec_memory_cache):
+    p = UPath(zipped_testdir_file, protocol="file")
+    t = p.move(UPath("memory:///myzipfile.zip"))
+    assert t.protocol == "memory"
+    assert t.exists()
+    yield t.as_uri()
+
+
+class TestChainedZipPath(TestZipPath):
+
+    @pytest.fixture(autouse=True)
+    def path(self, zipped_testdir_file_in_memory, request):
+        try:
+            (mode,) = request.param
+        except (ValueError, TypeError, AttributeError):
+            mode = "r"
+        self.path = UPath(
+            "zip://", fo="/myzipfile.zip", mode=mode, target_protocol="memory"
+        )

upath/tests/test_chain.py

@@ -5,6 +5,7 @@
 from fsspec.implementations.memory import MemoryFileSystem
 
 from upath import UPath
+from upath._chain import FSSpecChainParser
 
 
 @pytest.mark.parametrize(
@@ -115,3 +116,21 @@ def test_write_file(clear_memory_fs):
     pth = UPath("simplecache::memory://abc.txt")
     pth.write_bytes(b"hello world")
     assert clear_memory_fs.cat_file("abc.txt") == b"hello world"
+
+
+@pytest.mark.parametrize(
+    "urlpath",
+    [
+        "memory:///file.txt",
+        "simplecache::memory:///tmp",
+        "zip://file.txt::memory:///tmp.zip",
+        "zip://a/b/c.txt::simplecache::memory:///zipfile.zip",
+        "simplecache::zip://a/b/c.txt::tar://blah.zip::memory:///file.tar",
+    ],
+)
+def test_chain_parser_roundtrip(urlpath: str):
+    parser = FSSpecChainParser()
+    segments = parser.unchain(urlpath, {})
+    rechained, kw = parser.chain(segments)
+    assert rechained == urlpath
+    assert kw == {}