Handle fsspec.FSMap as zarr's store (use FSStore)

zarr-developers · Dec 20, 2022 · 3c173f9 · 3c173f9
1 parent 1af77b6
commit 3c173f9
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 0 deletions.
diff --git a/zarr/_storage/v3.py b/zarr/_storage/v3.py
@@ -567,6 +567,16 @@ def _normalize_store_arg_v3(store: Any, storage_options=None, mode="r") -> BaseS
         return store
     if isinstance(store, os.PathLike):
         store = os.fspath(store)
+    if FSStore._fsspec_installed():
+        import fsspec
+        if isinstance(store, fsspec.FSMap):
+            return FSStoreV3(store.root,
+                             fs=store.fs,
+                             mode=mode,
+                             check=store.check,
+                             create=store.create,
+                             missing_exceptions=store.missing_exceptions,
+                             **(storage_options or {}))
     if isinstance(store, str):
         if "://" in store or "::" in store:
             store = FSStoreV3(store, mode=mode, **(storage_options or {}))

diff --git a/zarr/storage.py b/zarr/storage.py
@@ -139,6 +139,16 @@ def _normalize_store_arg_v2(store: Any, storage_options=None, mode="r") -> BaseS
         return store
     if isinstance(store, os.PathLike):
         store = os.fspath(store)
+    if FSStore._fsspec_installed():
+        import fsspec
+        if isinstance(store, fsspec.FSMap):
+            return FSStore(store.root,
+                           fs=store.fs,
+                           mode=mode,
+                           check=store.check,
+                           create=store.create,
+                           missing_exceptions=store.missing_exceptions,
+                           **(storage_options or {}))
     if isinstance(store, str):
         if "://" in store or "::" in store:
             return FSStore(store, mode=mode, **(storage_options or {}))
@@ -1308,6 +1318,8 @@ def __init__(self, url, normalize_keys=False, key_separator=None,
                  create=False,
                  missing_exceptions=None,
                  **storage_options):
+        if not self._fsspec_installed():  # pragma: no cover
+            raise ImportError("`fsspec` is required to use zarr's FSStore")
         import fsspec
 
         mapper_options = {"check": check, "create": create}
@@ -1479,6 +1491,13 @@ def clear(self):
             raise ReadOnlyError()
         self.map.clear()
 
+    @classmethod
+    def _fsspec_installed(cls):
+        """Returns true if fsspec is installed"""
+        import importlib.util
+
+        return importlib.util.find_spec("fsspec") is not None
+
 
 class TempStore(DirectoryStore):
     """Directory store using a temporary directory for storage.

diff --git a/zarr/tests/test_storage.py b/zarr/tests/test_storage.py
@@ -2556,10 +2556,15 @@ def test_normalize_store_arg(tmpdir):
         assert isinstance(store, Class)
 
     if have_fsspec:
+        import fsspec
+
         path = tempfile.mkdtemp()
         store = normalize_store_arg("file://" + path, zarr_version=2, mode='w')
         assert isinstance(store, FSStore)
 
+        store = normalize_store_arg(fsspec.get_mapper("file://" + path))
+        assert isinstance(store, FSStore)
+
 
 def test_meta_prefix_6853():
 

diff --git a/zarr/tests/test_storage_v3.py b/zarr/tests/test_storage_v3.py
@@ -467,11 +467,16 @@ def test_normalize_store_arg_v3(tmpdir):
         normalize_store_arg(str(fn), zarr_version=3, mode='w', storage_options={"some": "kwargs"})
 
     if have_fsspec:
+        import fsspec
+
         path = tempfile.mkdtemp()
         store = normalize_store_arg("file://" + path, zarr_version=3, mode='w')
         assert isinstance(store, FSStoreV3)
         assert 'zarr.json' in store
 
+        store = normalize_store_arg(fsspec.get_mapper("file://" + path), zarr_version=3)
+        assert isinstance(store, FSStoreV3)
+
     fn = tmpdir.join('store.n5')
     with pytest.raises(NotImplementedError):
         normalize_store_arg(str(fn), zarr_version=3, mode='w')