fsspec 2024.9.0__py3-none-any.whl → 2024.12.0__py3-none-any.whl

fsspec/implementations/reference.py

@@ -5,9 +5,9 @@ import itertools
 import logging
 import math
 import os
-from itertools import chain
 from functools import lru_cache
-from typing import TYPE_CHECKING
+from itertools import chain
+from typing import TYPE_CHECKING, Literal
 
 import fsspec.core
 
@@ -20,6 +20,7 @@ except ImportError:
 from fsspec.asyn import AsyncFileSystem
 from fsspec.callbacks import DEFAULT_CALLBACK
 from fsspec.core import filesystem, open, split_protocol
+from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
 from fsspec.utils import isfilelike, merge_offset_ranges, other_paths
 
 logger = logging.getLogger("fsspec.reference")
@@ -41,7 +42,7 @@ def _first(d):
 
 def _prot_in_references(path, references):
     ref = references.get(path)
-    if isinstance(ref, (list, tuple)):
+    if isinstance(ref, (list, tuple)) and isinstance(ref[0], str):
         return split_protocol(ref[0])[0] if ref[0] else ref[0]
 
 
@@ -104,7 +105,13 @@ class LazyReferenceMapper(collections.abc.MutableMapping):
         return pd
 
     def __init__(
-        self, root, fs=None, out_root=None, cache_size=128, categorical_threshold=10
+        self,
+        root,
+        fs=None,
+        out_root=None,
+        cache_size=128,
+        categorical_threshold=10,
+        engine: Literal["fastparquet", "pyarrow"] = "fastparquet",
     ):
         """
 
@@ -126,16 +133,25 @@ class LazyReferenceMapper(collections.abc.MutableMapping):
             Encode urls as pandas.Categorical to reduce memory footprint if the ratio
             of the number of unique urls to total number of refs for each variable
             is greater than or equal to this number. (default 10)
+        engine: Literal["fastparquet","pyarrow"]
+            Engine choice for reading parquet files. (default is "fastparquet")
         """
+
         self.root = root
         self.chunk_sizes = {}
         self.out_root = out_root or self.root
         self.cat_thresh = categorical_threshold
+        self.engine = engine
         self.cache_size = cache_size
         self.url = self.root + "/{field}/refs.{record}.parq"
         # TODO: derive fs from `root`
         self.fs = fsspec.filesystem("file") if fs is None else fs
 
+        from importlib.util import find_spec
+
+        if self.engine == "pyarrow" and find_spec("pyarrow") is None:
+            raise ImportError("engine choice `pyarrow` is not installed.")
+
     def __getattr__(self, item):
         if item in ("_items", "record_size", "zmetadata"):
             self.setup()
@@ -158,8 +174,11 @@ class LazyReferenceMapper(collections.abc.MutableMapping):
             """cached parquet file loader"""
             path = self.url.format(field=field, record=record)
             data = io.BytesIO(self.fs.cat_file(path))
-            df = self.pd.read_parquet(data, engine="fastparquet")
-            refs = {c: df[c].to_numpy() for c in df.columns}
+            try:
+                df = self.pd.read_parquet(data, engine=self.engine)
+                refs = {c: df[c].to_numpy() for c in df.columns}
+            except OSError:
+                refs = None
             return refs
 
         self.open_refs = open_refs
@@ -413,7 +432,7 @@ class LazyReferenceMapper(collections.abc.MutableMapping):
         if len(partition) < self.record_size:
             try:
                 original = self.open_refs(field, record)
-            except IOError:
+            except OSError:
                 pass
 
         if original:
@@ -463,18 +482,28 @@ class LazyReferenceMapper(collections.abc.MutableMapping):
 
         fn = f"{base_url or self.out_root}/{field}/refs.{record}.parq"
         self.fs.mkdirs(f"{base_url or self.out_root}/{field}", exist_ok=True)
+
+        if self.engine == "pyarrow":
+            df_backend_kwargs = {"write_statistics": False}
+        elif self.engine == "fastparquet":
+            df_backend_kwargs = {
+                "stats": False,
+                "object_encoding": object_encoding,
+                "has_nulls": has_nulls,
+            }
+        else:
+            raise NotImplementedError(f"{self.engine} not supported")
+
         df.to_parquet(
             fn,
-            engine="fastparquet",
+            engine=self.engine,
             storage_options=storage_options
             or getattr(self.fs, "storage_options", None),
             compression="zstd",
             index=False,
-            stats=False,
-            object_encoding=object_encoding,
-            has_nulls=has_nulls,
-            # **kwargs,
+            **df_backend_kwargs,
         )
+
         partition.clear()
         self._items.pop((field, record))
 
@@ -486,6 +515,7 @@ class LazyReferenceMapper(collections.abc.MutableMapping):
         base_url: str
             Location of the output
         """
+
         # write what we have so far and clear sub chunks
         for thing in list(self._items):
             if isinstance(thing, tuple):
@@ -728,6 +758,10 @@ class ReferenceFileSystem(AsyncFileSystem):
             self.fss[remote_protocol] = fs
 
         self.fss[None] = fs or filesystem("file")  # default one
+        # Wrap any non-async filesystems to ensure async methods are available below
+        for k, f in self.fss.items():
+            if not f.async_impl:
+                self.fss[k] = AsyncFileSystemWrapper(f)
 
     def _cat_common(self, path, start=None, end=None):
         path = self._strip_protocol(path)
@@ -777,7 +811,9 @@ class ReferenceFileSystem(AsyncFileSystem):
             return part_or_url[start:end]
         protocol, _ = split_protocol(part_or_url)
         try:
-            await self.fss[protocol]._cat_file(part_or_url, start=start, end=end)
+            return await self.fss[protocol]._cat_file(
+                part_or_url, start=start0, end=end0
+            )
         except Exception as e:
             raise ReferenceNotReachable(path, part_or_url) from e
 
@@ -845,6 +881,9 @@ class ReferenceFileSystem(AsyncFileSystem):
                 # found and on_error is "raise"
                 try:
                     u, s, e = self._cat_common(p)
+                    if not isinstance(u, (bytes, str)):
+                        # nan/None from parquet
+                        continue
                 except FileNotFoundError as err:
                     if on_error == "raise":
                         raise
@@ -1147,13 +1186,17 @@ class ReferenceFileSystem(AsyncFileSystem):
         )  # ignores FileNotFound, just as well for directories
         self.dircache.clear()  # this is a bit heavy handed
 
-    async def _pipe_file(self, path, data):
+    async def _pipe_file(self, path, data, mode="overwrite", **kwargs):
+        if mode == "create" and self.exists(path):
+            raise FileExistsError
         # can be str or bytes
         self.references[path] = data
         self.dircache.clear()  # this is a bit heavy handed
 
-    async def _put_file(self, lpath, rpath, **kwargs):
+    async def _put_file(self, lpath, rpath, mode="overwrite", **kwargs):
         # puts binary
+        if mode == "create" and self.exists(rpath):
+            raise FileExistsError
         with open(lpath, "rb") as f:
             self.references[rpath] = f.read()
         self.dircache.clear()  # this is a bit heavy handed

fsspec/implementations/webhdfs.py

@@ -166,7 +166,8 @@ class WebHDFS(AbstractFileSystem):
             self.session.auth = HTTPBasicAuth(self.user, self.password)
 
     def _call(self, op, method="get", path=None, data=None, redirect=True, **kwargs):
-        url = self._apply_proxy(self.url + quote(path or "", safe="/="))
+        path = self._strip_protocol(path) if path is not None else ""
+        url = self._apply_proxy(self.url + quote(path, safe="/="))
         args = kwargs.copy()
         args.update(self.pars)
         args["op"] = op.upper()

fsspec/implementations/zip.py

@@ -1,3 +1,4 @@
+import os
 import zipfile
 
 import fsspec
@@ -48,7 +49,7 @@ class ZipFileSystem(AbstractArchiveFileSystem):
         if mode not in set("rwa"):
             raise ValueError(f"mode '{mode}' no understood")
         self.mode = mode
-        if isinstance(fo, str):
+        if isinstance(fo, (str, os.PathLike)):
             if mode == "a":
                 m = "r+b"
             else:

fsspec/mapping.py

@@ -112,7 +112,7 @@ class FSMap(MutableMapping):
             for k, v in out.items()
         }
         return {
-            key: out[k2]
+            key: out[k2] if on_error == "raise" else out.get(k2, KeyError(k2))
             for key, k2 in zip(keys, keys2)
             if on_error == "return" or not isinstance(out[k2], BaseException)
         }

fsspec/parquet.py

@@ -336,7 +336,7 @@ def _add_header_magic(data):
     # Add b"PAR1" to file headers
     for path in list(data.keys()):
         add_magic = True
-        for k in data[path].keys():
+        for k in data[path]:
             if k[0] == 0 and k[1] >= 4:
                 add_magic = False
                 break

fsspec/registry.py

@@ -202,6 +202,10 @@ known_implementations = {
         "err": 'SFTPFileSystem requires "paramiko" to be installed',
     },
     "tar": {"class": "fsspec.implementations.tar.TarFileSystem"},
+    "tosfs": {
+        "class": "tosfs.TosFileSystem",
+        "err": "Install tosfs to access ByteDance volcano engine Tinder Object Storage",
+    },
     "wandb": {"class": "wandbfs.WandbFS", "err": "Install wandbfs to access wandb"},
     "webdav": {
         "class": "webdav4.fsspec.WebdavFileSystem",

fsspec/spec.py

@@ -10,7 +10,7 @@ import weakref
 from errno import ESPIPE
 from glob import has_magic
 from hashlib import sha256
-from typing import Any, ClassVar, Dict, Tuple
+from typing import Any, ClassVar
 
 from .callbacks import DEFAULT_CALLBACK
 from .config import apply_config, conf
@@ -117,8 +117,8 @@ class AbstractFileSystem(metaclass=_Cached):
     _extra_tokenize_attributes = ()
 
     # Set by _Cached metaclass
-    storage_args: Tuple[Any, ...]
-    storage_options: Dict[str, Any]
+    storage_args: tuple[Any, ...]
+    storage_options: dict[str, Any]
 
     def __init__(self, *args, **storage_options):
         """Create and configure file-system instance
@@ -408,7 +408,7 @@ class AbstractFileSystem(metaclass=_Cached):
         topdown: bool (True)
             Whether to walk the directory tree from the top downwards or from
             the bottom upwards.
-        on_error: "omit", "raise", a collable
+        on_error: "omit", "raise", a callable
             if omit (default), path with exception will simply be empty;
             If raise, an underlying exception will be raised;
             if callable, it will be called with a single OSError instance as argument
@@ -428,11 +428,9 @@ class AbstractFileSystem(metaclass=_Cached):
         except (FileNotFoundError, OSError) as e:
             if on_error == "raise":
                 raise
-            elif callable(on_error):
+            if callable(on_error):
                 on_error(e)
-            if detail:
-                return path, {}, {}
-            return path, [], []
+            return
 
         for info in listing:
             # each info name must be at least [path]/part , but here
@@ -617,11 +615,9 @@ class AbstractFileSystem(metaclass=_Cached):
             p: info
             for p, info in sorted(allpaths.items())
             if pattern.match(
-                (
-                    p + "/"
-                    if append_slash_to_dirname and info["type"] == "directory"
-                    else p
-                )
+                p + "/"
+                if append_slash_to_dirname and info["type"] == "directory"
+                else p
             )
         }
 
@@ -650,7 +646,7 @@ class AbstractFileSystem(metaclass=_Cached):
         Returns a single dictionary, with exactly the same information as ``ls``
         would with ``detail=True``.
 
-        The default implementation should calls ls and could be overridden by a
+        The default implementation calls ls and could be overridden by a
         shortcut. kwargs are passed on to ```ls()``.
 
         Some file systems might not be able to measure the file's size, in
@@ -782,8 +778,12 @@ class AbstractFileSystem(metaclass=_Cached):
                 return f.read(end - f.tell())
             return f.read()
 
-    def pipe_file(self, path, value, **kwargs):
+    def pipe_file(self, path, value, mode="overwrite", **kwargs):
         """Set the bytes of given file"""
+        if mode == "create" and self.exists(path):
+            # non-atomic but simple way; or could use "xb" in open(), which is likely
+            # not as well supported
+            raise FileExistsError
         with self.open(path, "wb", **kwargs) as f:
             f.write(value)
 
@@ -975,8 +975,12 @@ class AbstractFileSystem(metaclass=_Cached):
             with callback.branched(rpath, lpath) as child:
                 self.get_file(rpath, lpath, callback=child, **kwargs)
 
-    def put_file(self, lpath, rpath, callback=DEFAULT_CALLBACK, **kwargs):
+    def put_file(
+        self, lpath, rpath, callback=DEFAULT_CALLBACK, mode="overwrite", **kwargs
+    ):
         """Copy single file to remote"""
+        if mode == "create" and self.exists(rpath):
+            raise FileExistsError
         if os.path.isdir(lpath):
             self.makedirs(rpath, exist_ok=True)
             return None
@@ -1266,6 +1270,9 @@ class AbstractFileSystem(metaclass=_Cached):
             Target file
         mode: str like 'rb', 'w'
             See builtin ``open()``
+            Mode "x" (exclusive write) may be implemented by the backend. Even if
+            it is, whether  it is checked up front or on commit, and whether it is
+            atomic is implementation-dependent.
         block_size: int
             Some indication of buffering - this is a value in bytes
         cache_options : dict, optional
@@ -1444,7 +1451,7 @@ class AbstractFileSystem(metaclass=_Cached):
 
         return json.loads(blob, cls=FilesystemJSONDecoder)
 
-    def to_dict(self, *, include_password: bool = True) -> Dict[str, Any]:
+    def to_dict(self, *, include_password: bool = True) -> dict[str, Any]:
         """
         JSON-serializable dictionary representation of this filesystem instance.
 
@@ -1485,7 +1492,7 @@ class AbstractFileSystem(metaclass=_Cached):
         )
 
     @staticmethod
-    def from_dict(dct: Dict[str, Any]) -> AbstractFileSystem:
+    def from_dict(dct: dict[str, Any]) -> AbstractFileSystem:
         """
         Recreate a filesystem instance from dictionary representation.
 
@@ -1569,6 +1576,141 @@ class AbstractFileSystem(metaclass=_Cached):
         """Return the modified timestamp of a file as a datetime.datetime"""
         raise NotImplementedError
 
+    def tree(
+        self,
+        path: str = "/",
+        recursion_limit: int = 2,
+        max_display: int = 25,
+        display_size: bool = False,
+        prefix: str = "",
+        is_last: bool = True,
+        first: bool = True,
+        indent_size: int = 4,
+    ) -> str:
+        """
+        Return a tree-like structure of the filesystem starting from the given path as a string.
+
+        Parameters
+        ----------
+            path: Root path to start traversal from
+            recursion_limit: Maximum depth of directory traversal
+            max_display: Maximum number of items to display per directory
+            display_size: Whether to display file sizes
+            prefix: Current line prefix for visual tree structure
+            is_last: Whether current item is last in its level
+            first: Whether this is the first call (displays root path)
+            indent_size: Number of spaces by indent
+
+        Returns
+        -------
+            str: A string representing the tree structure.
+
+        Example
+        -------
+            >>> from fsspec import filesystem
+
+            >>> fs = filesystem('ftp', host='test.rebex.net', user='demo', password='password')
+            >>> tree = fs.tree(display_size=True, recursion_limit=3, indent_size=8, max_display=10)
+            >>> print(tree)
+        """
+
+        def format_bytes(n: int) -> str:
+            """Format bytes as text."""
+            for prefix, k in (
+                ("P", 2**50),
+                ("T", 2**40),
+                ("G", 2**30),
+                ("M", 2**20),
+                ("k", 2**10),
+            ):
+                if n >= 0.9 * k:
+                    return f"{n / k:.2f} {prefix}b"
+            return f"{n}B"
+
+        result = []
+
+        if first:
+            result.append(path)
+
+        if recursion_limit:
+            indent = " " * indent_size
+            contents = self.ls(path, detail=True)
+            contents.sort(
+                key=lambda x: (x.get("type") != "directory", x.get("name", ""))
+            )
+
+            if max_display is not None and len(contents) > max_display:
+                displayed_contents = contents[:max_display]
+                remaining_count = len(contents) - max_display
+            else:
+                displayed_contents = contents
+                remaining_count = 0
+
+            for i, item in enumerate(displayed_contents):
+                is_last_item = (i == len(displayed_contents) - 1) and (
+                    remaining_count == 0
+                )
+
+                branch = (
+                    "└" + ("─" * (indent_size - 2))
+                    if is_last_item
+                    else "├" + ("─" * (indent_size - 2))
+                )
+                branch += " "
+                new_prefix = prefix + (
+                    indent if is_last_item else "│" + " " * (indent_size - 1)
+                )
+
+                name = os.path.basename(item.get("name", ""))
+
+                if display_size and item.get("type") == "directory":
+                    sub_contents = self.ls(item.get("name", ""), detail=True)
+                    num_files = sum(
+                        1 for sub_item in sub_contents if sub_item.get("type") == "file"
+                    )
+                    num_folders = sum(
+                        1
+                        for sub_item in sub_contents
+                        if sub_item.get("type") == "directory"
+                    )
+
+                    if num_files == 0 and num_folders == 0:
+                        size = " (empty folder)"
+                    elif num_files == 0:
+                        size = f" ({num_folders} subfolder{'s' if num_folders > 1 else ''})"
+                    elif num_folders == 0:
+                        size = f" ({num_files} file{'s' if num_files > 1 else ''})"
+                    else:
+                        size = f" ({num_files} file{'s' if num_files > 1 else ''}, {num_folders} subfolder{'s' if num_folders > 1 else ''})"
+                elif display_size and item.get("type") == "file":
+                    size = f" ({format_bytes(item.get('size', 0))})"
+                else:
+                    size = ""
+
+                result.append(f"{prefix}{branch}{name}{size}")
+
+                if item.get("type") == "directory" and recursion_limit > 0:
+                    result.append(
+                        self.tree(
+                            path=item.get("name", ""),
+                            recursion_limit=recursion_limit - 1,
+                            max_display=max_display,
+                            display_size=display_size,
+                            prefix=new_prefix,
+                            is_last=is_last_item,
+                            first=False,
+                            indent_size=indent_size,
+                        )
+                    )
+
+            if remaining_count > 0:
+                more_message = f"{remaining_count} more item(s) not displayed."
+                result.append(
+                    f"{prefix}{'└' + ('─' * (indent_size - 2))} {more_message}"
+                )
+
+        return "\n".join(_ for _ in result if _)
+
     # ------------------------------------------------------------------------
     # Aliases
 
@@ -1733,7 +1875,7 @@ class AbstractBufferedFile(io.IOBase):
 
         self.kwargs = kwargs
 
-        if mode not in {"ab", "rb", "wb"}:
+        if mode not in {"ab", "rb", "wb", "xb"}:
             raise NotImplementedError("File mode not supported")
         if mode == "rb":
             if size is not None:
@@ -1799,7 +1941,7 @@ class AbstractBufferedFile(io.IOBase):
 
     def info(self):
         """File information about this path"""
-        if "r" in self.mode:
+        if self.readable():
             return self.details
         else:
             raise ValueError("Info not available while writing")
@@ -1846,7 +1988,7 @@ class AbstractBufferedFile(io.IOBase):
         data: bytes
             Set of bytes to be written.
         """
-        if self.mode not in {"wb", "ab"}:
+        if not self.writable():
             raise ValueError("File not in write mode")
         if self.closed:
             raise ValueError("I/O operation on closed file.")
@@ -1879,7 +2021,7 @@ class AbstractBufferedFile(io.IOBase):
         if force:
             self.forced = True
 
-        if self.mode not in {"wb", "ab"}:
+        if self.readable():
             # no-op to flush on read-mode
             return
 
@@ -1917,7 +2059,7 @@ class AbstractBufferedFile(io.IOBase):
 
     def _fetch_range(self, start, end):
         """Get the specified set of bytes from remote"""
-        raise NotImplementedError
+        return self.fs.cat_file(self.path, start=start, end=end)
 
     def read(self, length=-1):
         """
@@ -2028,21 +2170,22 @@ class AbstractBufferedFile(io.IOBase):
             return
         if self.closed:
             return
-        if self.mode == "rb":
-            self.cache = None
-        else:
-            if not self.forced:
-                self.flush(force=True)
-
-            if self.fs is not None:
-                self.fs.invalidate_cache(self.path)
-                self.fs.invalidate_cache(self.fs._parent(self.path))
+        try:
+            if self.mode == "rb":
+                self.cache = None
+            else:
+                if not self.forced:
+                    self.flush(force=True)
 
-        self.closed = True
+                if self.fs is not None:
+                    self.fs.invalidate_cache(self.path)
+                    self.fs.invalidate_cache(self.fs._parent(self.path))
+        finally:
+            self.closed = True
 
     def readable(self):
         """Whether opened for reading"""
-        return self.mode == "rb" and not self.closed
+        return "r" in self.mode and not self.closed
 
     def seekable(self):
         """Whether is seekable (only in read mode)"""
@@ -2050,7 +2193,23 @@ class AbstractBufferedFile(io.IOBase):
 
     def writable(self):
         """Whether opened for writing"""
-        return self.mode in {"wb", "ab"} and not self.closed
+        return self.mode in {"wb", "ab", "xb"} and not self.closed
+
+    def __reduce__(self):
+        if self.mode != "rb":
+            raise RuntimeError("Pickling a writeable file is not supported")
+
+        return reopen, (
+            self.fs,
+            self.path,
+            self.mode,
+            self.blocksize,
+            self.loc,
+            self.size,
+            self.autocommit,
+            self.cache.name if self.cache else "none",
+            self.kwargs,
+        )
 
     def __del__(self):
         if not self.closed:
@@ -2066,3 +2225,18 @@ class AbstractBufferedFile(io.IOBase):
 
     def __exit__(self, *args):
         self.close()
+
+
+def reopen(fs, path, mode, blocksize, loc, size, autocommit, cache_type, kwargs):
+    file = fs.open(
+        path,
+        mode=mode,
+        block_size=blocksize,
+        autocommit=autocommit,
+        cache_type=cache_type,
+        size=size,
+        **kwargs,
+    )
+    if loc > 0:
+        file.seek(loc)
+    return file

fsspec/tests/abstract/__init__.py

@@ -6,6 +6,8 @@ import pytest
 from fsspec.implementations.local import LocalFileSystem
 from fsspec.tests.abstract.copy import AbstractCopyTests  # noqa: F401
 from fsspec.tests.abstract.get import AbstractGetTests  # noqa: F401
+from fsspec.tests.abstract.open import AbstractOpenTests  # noqa: F401
+from fsspec.tests.abstract.pipe import AbstractPipeTests  # noqa: F401
 from fsspec.tests.abstract.put import AbstractPutTests  # noqa: F401
 
 
@@ -225,7 +227,7 @@ class BaseAbstractFixtures:
         for i in range(10):
             hashed_i = md5(str(i).encode("utf-8")).hexdigest()
             path = some_join(source, f"{hashed_i}.txt")
-            some_fs.pipe(path=path, value=f"{i}".encode("utf-8"))
+            some_fs.pipe(path=path, value=f"{i}".encode())
         return source
 
 

fsspec/tests/abstract/open.py

@@ -0,0 +1,11 @@
+import pytest
+
+
+class AbstractOpenTests:
+    def test_open_exclusive(self, fs, fs_target):
+        with fs.open(fs_target, "wb") as f:
+            f.write(b"data")
+        with fs.open(fs_target, "rb") as f:
+            assert f.read() == b"data"
+        with pytest.raises(FileExistsError):
+            fs.open(fs_target, "xb")

fsspec/tests/abstract/pipe.py

@@ -0,0 +1,11 @@
+import pytest
+
+
+class AbstractPipeTests:
+    def test_pipe_exclusive(self, fs, fs_target):
+        fs.pipe_file(fs_target, b"data")
+        assert fs.cat_file(fs_target) == b"data"
+        with pytest.raises(FileExistsError):
+            fs.pipe_file(fs_target, b"data", mode="create")
+        fs.pipe_file(fs_target, b"new data", mode="overwrite")
+        assert fs.cat_file(fs_target) == b"new data"