From 3c173f977d92348df1a36b139238de3464d2b8e5 Mon Sep 17 00:00:00 2001
From: Rafal Wojdyla <ravwojdyla@gmail.com>
Date: Tue, 20 Dec 2022 22:17:33 +0900
Subject: [PATCH] Handle fsspec.FSMap as zarr's store (use FSStore)

---
 zarr/_storage/v3.py           | 10 ++++++++++
 zarr/storage.py               | 19 +++++++++++++++++++
 zarr/tests/test_storage.py    |  5 +++++
 zarr/tests/test_storage_v3.py |  5 +++++
 4 files changed, 39 insertions(+)

diff --git a/zarr/_storage/v3.py b/zarr/_storage/v3.py
index 515e6f5aaa..a0a1870ffc 100644
--- 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
index 4acf637330..a2a8919d0b 100644
--- 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
index 39d4b5988d..7c23735f36 100644
--- 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
index 13b5011676..4f6215135c 100644
--- 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')