from __future__ import annotations
import fsspec
from fsspec.config import conf
from fsspec.spec import AbstractBufferedFile
from fsspec_rs.fsspec_rs import RustS3Fs
[docs]
class S3FileSystem(fsspec.AbstractFileSystem):
"""S3 filesystem backed by the Rust ``S3Fs`` implementation.
The protocol is registered as ``"s3-rs"`` so it can coexist with the
pure-Python ``s3fs.S3FileSystem`` which owns the ``"s3"`` protocol.
Credential resolution (matching s3fs conventions):
1. Explicit keyword arguments (``key``, ``secret``, ``endpoint_url``, …)
2. Values from ``fsspec.config.conf["s3"]`` (populated from ``FSSPEC_S3_*``
environment variables)
3. Falls through to object_store's own credential chain (IAM roles, etc.)
"""
protocol = ("s3-rs",)
@classmethod
def _strip_protocol(cls, path):
"""Strip ``s3-rs://`` prefix, returning ``bucket/key``."""
for proto in cls.protocol:
prefix = f"{proto}://"
if path.startswith(prefix):
return path[len(prefix) :].lstrip("/")
return path
@staticmethod
def _get_kwargs_from_urls(path):
"""Extract ``bucket`` from a ``s3-rs://bucket/key`` URL."""
for proto in ("s3-rs",):
prefix = f"{proto}://"
if path.startswith(prefix):
rest = path[len(prefix) :].lstrip("/")
bucket = rest.split("/", 1)[0]
if bucket:
return {"bucket": bucket}
return {}
[docs]
def unstrip_protocol(self, path):
"""Reconstruct ``s3-rs://bucket/key`` from ``bucket/key``."""
bucket_prefix = f"{self.bucket}/"
if path.startswith(bucket_prefix):
key = path[len(bucket_prefix) :]
return f"s3-rs://{self.bucket}/{key}" if key else f"s3-rs://{self.bucket}"
return f"s3-rs://{self.bucket}/{path}" if path else f"s3-rs://{self.bucket}"
def __init__(
self,
bucket: str | None = None,
key: str | None = None,
secret: str | None = None,
endpoint_url: str | None = None,
region: str | None = None,
token: str | None = None,
anon: bool = False,
client_kwargs: dict | None = None,
**storage_options,
):
super().__init__(**storage_options)
if not bucket:
raise ValueError("bucket is required unless it is provided by an s3-rs:// URL")
# Merge from fsspec.config.conf (populated by FSSPEC_S3_* env vars)
s3_conf = conf.get("s3", {})
key = key or s3_conf.get("key")
secret = secret or s3_conf.get("secret")
endpoint_url = endpoint_url or s3_conf.get("endpoint_url")
region = region or s3_conf.get("region")
token = token or s3_conf.get("token")
# Also honour client_kwargs.endpoint_url (s3fs convention)
if client_kwargs and not endpoint_url:
endpoint_url = client_kwargs.get("endpoint_url")
self._rust = RustS3Fs(
bucket=bucket,
key=key,
secret=secret,
endpoint_url=endpoint_url,
region=region,
token=token,
anon=anon,
)
self.bucket = bucket
# Core primitives — delegated to Rust
[docs]
def ls(self, path: str, detail: bool = True, **kwargs):
return self._rust.ls(path, detail=detail)
[docs]
def info(self, path: str, **kwargs):
return self._rust.info(path)
def _open(
self,
path: str,
mode: str = "rb",
block_size=None,
autocommit: bool = True,
cache_options=None,
cache_type=None,
**kwargs,
):
"""Return a Rust-backed S3 file object."""
return S3File(
self,
path,
mode=mode,
cache_type=cache_type,
block_size=block_size,
)
[docs]
def mkdir(self, path: str, create_parents: bool = True, **kwargs):
return self._rust.mkdir(path, create_parents=create_parents)
[docs]
def rmdir(self, path: str):
return self._rust.rmdir(path)
[docs]
def rm_file(self, path: str):
return self._rust.rm_file(path)
[docs]
def rm(self, path, recursive: bool = False, maxdepth=None):
if isinstance(path, str):
return self._rust.rm(path, recursive=recursive)
for p in path:
self._rust.rm(p, recursive=recursive)
[docs]
def cp_file(self, path1: str, path2: str, **kwargs):
return self._rust.cp_file(path1, path2)
# Higher-level helpers — delegated to Rust for speed
[docs]
def exists(self, path: str, **kwargs) -> bool:
return self._rust.exists(path)
[docs]
def isdir(self, path: str) -> bool:
return self._rust.isdir(path)
[docs]
def isfile(self, path: str) -> bool:
return self._rust.isfile(path)
[docs]
def size(self, path: str) -> int:
return self._rust.size(path)
[docs]
def cat_file(self, path: str, start=None, end=None, **kwargs) -> bytes:
return bytes(self._rust.cat_file(path, start=start, end=end))
[docs]
def pipe_file(self, path: str, value: bytes, **kwargs):
return self._rust.pipe_file(path, value)
[docs]
def head(self, path: str, size: int = 1024) -> bytes:
return bytes(self._rust.head(path, size))
[docs]
def tail(self, path: str, size: int = 1024) -> bytes:
return bytes(self._rust.tail(path, size))
[docs]
def walk(self, path: str, maxdepth=None, topdown: bool = True, **kwargs):
entries = self._rust.walk(path, max_depth=maxdepth, topdown=topdown)
yield from entries
[docs]
def find(self, path: str, maxdepth=None, withdirs: bool = False, **kwargs):
return self._rust.find(path, max_depth=maxdepth, with_dirs=withdirs)
[docs]
def read_text(self, path: str, encoding=None, errors=None, newline=None, **kwargs) -> str:
return self._rust.read_text(path)
def __repr__(self) -> str:
return f"S3FileSystem(bucket='{self.bucket}')"
[docs]
class S3File(AbstractBufferedFile):
"""File-like wrapper delegating to the Rust ``RustS3File``."""
def __init__(self, fs: S3FileSystem, path: str, mode: str = "rb", cache_type=None, block_size=None, **kwargs):
self._rust_file = None # set early so __del__/closed don't blow up
open_kwargs = {}
if cache_type is not None:
open_kwargs["cache_type"] = cache_type
if block_size is not None:
open_kwargs["block_size"] = block_size
self._rust_file = fs._rust.open(path, mode, **open_kwargs)
self.path = path
self.mode = mode
self.fs = fs
# Minimal attributes — don't call super().__init__()
self.blocksize = 0
self.loc = 0
# io.RawIOBase-like interface
[docs]
def read(self, length: int = -1) -> bytes:
data = bytes(self._rust_file.read(length))
self.loc += len(data)
return data
[docs]
def write(self, data: bytes) -> int:
n = self._rust_file.write(data)
self.loc += n
return n
[docs]
def seek(self, loc: int, whence: int = 0) -> int:
pos = self._rust_file.seek(loc, whence)
self.loc = pos
return pos
[docs]
def tell(self) -> int:
return self._rust_file.tell()
[docs]
def flush(self, force: bool = False):
self._rust_file.flush()
[docs]
def close(self):
self._rust_file.close()
@property
def closed(self) -> bool:
return self._rust_file is None or self._rust_file.closed
[docs]
def readable(self) -> bool:
return "r" in self.mode
[docs]
def writable(self) -> bool:
return "w" in self.mode or "a" in self.mode or "x" in self.mode
[docs]
def seekable(self) -> bool:
return True
def __enter__(self):
return self
def __exit__(self, *args):
self.close()
def __repr__(self) -> str:
state = "closed" if self.closed else "open"
return f"S3File('{self.path}', {state})"