Skip to content

Commit

Permalink
fsspec: Allow calling register with no arguments (#298)
Browse files Browse the repository at this point in the history
* Allow calling `register` with no arguments

* Don't register as default for `file://` and `memory://`
  • Loading branch information
kylebarron authored Mar 3, 2025
1 parent 1e3162b commit 94b69cc
Showing 1 changed file with 49 additions and 8 deletions.
57 changes: 49 additions & 8 deletions obstore/python/obstore/fsspec.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@
S3ConfigInput,
)

SUPPORTED_PROTOCOLS = {
SUPPORTED_PROTOCOLS: set[str] = {
"abfs",
"abfss",
"adl",
Expand All @@ -80,6 +80,24 @@
"s3",
"s3a",
}
"""All supported protocols."""

SUPPORTED_PROTOCOLS_T = Literal[
"abfs",
"abfss",
"adl",
"az",
"azure",
"file",
"gcs",
"gs",
"http",
"https",
"memory",
"s3",
"s3a",
]
"""A type hint for all supported protocols."""


class FsspecStore(fsspec.asyn.AsyncFileSystem):
Expand Down Expand Up @@ -747,24 +765,44 @@ def loc(self, value: int) -> None:
raise ValueError("Cannot set `.loc`. Use `seek` instead.")


def register(protocol: str | Iterable[str], *, asynchronous: bool = False) -> None:
def register(
protocol: SUPPORTED_PROTOCOLS_T
| str
| Iterable[SUPPORTED_PROTOCOLS_T]
| Iterable[str]
| None = None,
*,
asynchronous: bool = False,
) -> None:
"""Dynamically register a subclass of FsspecStore for the given protocol(s).
This function creates a new subclass of FsspecStore with the specified
protocol and registers it with fsspec. If multiple protocols are provided,
the function registers each one individually.
Args:
protocol (str | list[str]): A single protocol (e.g., "s3", "gcs", "abfs") or
a list of protocols to register FsspecStore for.
asynchronous (bool, optional): If True, the registered store will support
asynchronous operations. Defaults to False.
protocol: A single protocol (e.g., "s3", "gcs", "abfs") or
a list of protocols to register FsspecStore for. Defaults to `None`, which
will register `obstore` as the provider for all [supported
protocols][obstore.fsspec.SUPPORTED_PROTOCOLS] **except** for `file://` and
`memory://`. If you wish to use `obstore` via fsspec for `file://` or
`memory://` URLs, list them explicitly.
asynchronous: If `True`, the registered store will support
asynchronous operations. Defaults to `False`.
Example:
```py
# Register obstore as the default handler for all supported protocols except for
# `memory://` and `file://`
register()
register("s3")
register("s3", asynchronous=True) # Registers an async store for "s3"
register(["gcs", "abfs"]) # Registers both "gcs" and "abfs"
# Registers an async store for "s3"
register("s3", asynchronous=True)
# Registers both "gcs" and "abfs"
register(["gcs", "abfs"])
```
Notes:
Expand All @@ -773,6 +811,9 @@ def register(protocol: str | Iterable[str], *, asynchronous: bool = False) -> No
FsspecStore class.
"""
if protocol is None:
protocol = SUPPORTED_PROTOCOLS - {"file", "memory"}

if isinstance(protocol, str):
_register(protocol, asynchronous=asynchronous)
return
Expand Down

0 comments on commit 94b69cc

Please sign in to comment.