WIP: Support fsspec mutable mapping objects in zarr.open

Author: maxrjones

src/zarr/storage/_common.py

@@ -312,8 +312,27 @@ async def make_store_path(
             # By only allowing dictionaries, which are in-memory, we know that MemoryStore appropriate.
             store = await MemoryStore.open(store_dict=store_like, read_only=_read_only)
         else:
-            msg = f"Unsupported type for store_like: '{type(store_like).__name__}'"  # type: ignore[unreachable]
-            raise TypeError(msg)
+            try:  # type: ignore[unreachable]
+                import fsspec
+
+                if isinstance(store_like, fsspec.mapping.FSMap):
+                    if path:
+                        raise TypeError(
+                            "'path' was provided but is not used for FSMap store_like objects"
+                        )
+                    if storage_options:
+                        raise TypeError(
+                            "'storage_options was provided but is not used for FSMap store_like objects"
+                        )
+                    store = FsspecStore.from_mapper(store_like, read_only=_read_only)
+                else:
+                    raise (
+                        TypeError(f"Unsupported type for store_like: '{type(store_like).__name__}'")
+                    )
+            except ImportError:
+                raise (
+                    TypeError(f"Unsupported type for store_like: '{type(store_like).__name__}'")
+                ) from None
 
         result = await StorePath.open(store, path=path_normalized, mode=mode)
 

src/zarr/storage/_fsspec.py

@@ -16,7 +16,9 @@
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterable
 
+    from fsspec import AbstractFileSystem
     from fsspec.asyn import AsyncFileSystem
+    from fsspec.mapping import FSMap
 
     from zarr.core.buffer import BufferPrototype
     from zarr.core.common import BytesLike
@@ -29,6 +31,20 @@
 )
 
 
+def _make_async(fs: AbstractFileSystem) -> AsyncFileSystem:
+    try:
+        from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
+
+        fs = AsyncFileSystemWrapper(fs)
+    except ImportError as e:
+        raise ImportError(
+            f"The filesystem '{fs}' is synchronous, and the required "
+            "AsyncFileSystemWrapper is not available. Upgrade fsspec to version "
+            "2024.12.0 or later to enable this functionality."
+        ) from e
+    return fs
+
+
 class FsspecStore(Store):
     """
     A remote Store based on FSSpec
@@ -136,6 +152,37 @@ def from_upath(
             allowed_exceptions=allowed_exceptions,
         )
 
+    @classmethod
+    def from_mapper(
+        cls,
+        fs_map: FSMap,
+        read_only: bool = False,
+        allowed_exceptions: tuple[type[Exception], ...] = ALLOWED_EXCEPTIONS,
+    ) -> FsspecStore:
+        """
+        Create a FsspecStore from an upath object.
+
+        Parameters
+        ----------
+        read_only : bool
+            Whether the store is read-only, defaults to False.
+        allowed_exceptions : tuple, optional
+            The exceptions that are allowed to be raised when accessing the
+            store. Defaults to ALLOWED_EXCEPTIONS.
+
+        Returns
+        -------
+        FsspecStore
+        """
+        if not fs_map.fs.async_impl or not fs_map.fs.asynchronous:
+            raise TypeError("Filesystem needs to support async operations.")
+        return cls(
+            fs=fs_map.fs,
+            path=fs_map.root,
+            read_only=read_only,
+            allowed_exceptions=allowed_exceptions,
+        )
+
     @classmethod
     def from_url(
         cls,
@@ -174,16 +221,7 @@ def from_url(
 
         fs, path = url_to_fs(url, **opts)
         if not fs.async_impl:
-            try:
-                from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
-
-                fs = AsyncFileSystemWrapper(fs)
-            except ImportError as e:
-                raise ImportError(
-                    f"The filesystem for URL '{url}' is synchronous, and the required "
-                    "AsyncFileSystemWrapper is not available. Upgrade fsspec to version "
-                    "2024.12.0 or later to enable this functionality."
-                ) from e
+            fs = _make_async(fs)
 
         # fsspec is not consistent about removing the scheme from the path, so check and strip it here
         # https://github.com/fsspec/filesystem_spec/issues/1722

tests/test_api.py

@@ -288,6 +288,23 @@ def test_open_with_mode_w_minus(tmp_path: pathlib.Path) -> None:
         zarr.open(store=tmp_path, mode="w-")
 
 
+@pytest.mark.xfail(
+    reason="Automatic sync -> async filesystems not implemented yet for FSMap objects."
+)
+def test_open_fsmap_file(tmp_path: pathlib.Path) -> None:
+    fsspec = pytest.importorskip("fsspec")
+    fs = fsspec.filesystem("file")
+    mapper = fs.get_mapper(tmp_path)
+    arr = zarr.open(store=mapper, mode="w", shape=(3, 3))
+    assert isinstance(arr, Array)
+
+    arr[...] = 3
+    z2 = zarr.open(store=mapper, mode="w", shape=(3, 3))
+    assert isinstance(z2, Array)
+    assert not (z2[:] == 3).all()
+    z2[:] = 3
+
+
 @pytest.mark.parametrize("zarr_format", [2, 3])
 def test_array_order(zarr_format: ZarrFormat) -> None:
     arr = zarr.ones(shape=(2, 2), order=None, zarr_format=zarr_format)

tests/test_store/test_fsspec.py

@@ -9,6 +9,7 @@
 from packaging.version import parse as parse_version
 
 import zarr.api.asynchronous
+from zarr import Array
 from zarr.abc.store import OffsetByteRequest
 from zarr.core.buffer import Buffer, cpu, default_buffer_prototype
 from zarr.core.sync import _collect_aiterator, sync
@@ -104,6 +105,19 @@ async def test_basic() -> None:
     assert out[0].to_bytes() == data[1:]
 
 
+def test_open_s3map() -> None:
+    s3_filesystem = s3fs.S3FileSystem(asynchronous=True, endpoint_url=endpoint_url, anon=False)
+    mapper = s3_filesystem.get_mapper(f"s3://{test_bucket_name}/map/foo/")
+    arr = zarr.open(store=mapper, mode="w", shape=(3, 3))
+    assert isinstance(arr, Array)
+
+    arr[...] = 3
+    z2 = zarr.open(store=mapper, mode="w", shape=(3, 3))
+    assert isinstance(z2, Array)
+    assert not (z2[:] == 3).all()
+    z2[:] = 3
+
+
 class TestFsspecStoreS3(StoreTests[FsspecStore, cpu.Buffer]):
     store_cls = FsspecStore
     buffer_cls = cpu.Buffer