From 453811c5209f3115ed8e814ef4cd626adb227450 Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Mon, 27 Jan 2025 08:29:44 +0100
Subject: [PATCH] First check for compatibility for zarr version 3 in storage
 backend

---
 modelrunner/storage/backend/zarr.py | 22 ++++++++++++++++++----
 tests/helpers/storage.py            |  6 ++++--
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/modelrunner/storage/backend/zarr.py b/modelrunner/storage/backend/zarr.py
index bfa0415..c7ef100 100644
--- a/modelrunner/storage/backend/zarr.py
+++ b/modelrunner/storage/backend/zarr.py
@@ -14,7 +14,18 @@
 import numpy as np
 import zarr
 from numpy.typing import ArrayLike, DTypeLike
-from zarr._storage.store import Store
+
+try:
+    # try import path from zarr version 3
+    from zarr.abc.store import Store
+
+    is_zarr2 = False
+
+except ImportError:
+    # try import path from zarr version 2
+    from zarr._storage.store import Store
+
+    is_zarr2 = True
 
 from ..access_modes import ModeType
 from ..attributes import AttrsLike
@@ -55,7 +66,10 @@ def __init__(self, store_or_path: str | Path | Store, *, mode: ModeType = "read"
                 if self.mode.file_mode == "r":
                     self._logger.info("DirectoryStore is always opened writable")
 
-                self._store = zarr.DirectoryStore(path)
+                if is_zarr2:
+                    self._store = zarr.DirectoryStore(path)
+                else:
+                    self._store = zarr.LocalStore(path)
 
             elif path.suffix == ".zip":
                 # create a ZipStore
@@ -65,7 +79,7 @@ def __init__(self, store_or_path: str | Path | Store, *, mode: ModeType = "read"
                     path.unlink()
                 self._store = zarr.storage.ZipStore(path, mode=file_mode)
 
-            elif path.suffix == ".sqldb":
+            elif is_zarr2 and path.suffix == ".sqldb":
                 # create a SQLiteStore
                 if self.mode.file_mode == "w" and path.exists():
                     self._logger.info("Delete file `%s`", path)
@@ -135,7 +149,7 @@ def keys(self, loc: Sequence[str] | None = None) -> Collection[str]:
             return self._root.keys()  # type: ignore
 
     def is_group(self, loc: Sequence[str], *, ignore_cls: bool = False) -> bool:
-        return isinstance(self[loc], zarr.hierarchy.Group)
+        return isinstance(self[loc], zarr.Group)
 
     def _create_group(self, loc: Sequence[str]) -> None:
         parent, name = self._get_parent(loc)
diff --git a/tests/helpers/storage.py b/tests/helpers/storage.py
index 263c179..50f77d8 100644
--- a/tests/helpers/storage.py
+++ b/tests/helpers/storage.py
@@ -42,7 +42,7 @@ def storage_extensions(
         dot (bool):
             Indicates whether the returned extensions are prepended with a dot (`.`)
         exclude (list, optional):
-            Extensions (without dots) that should be explicitely excluded
+            Extensions (without dots) that should be explicitly excluded
 
     Returns:
         sequence: sorted list of extensions
@@ -53,10 +53,12 @@ def storage_extensions(
     if module_available("h5py"):
         exts.append("hdf")
     if module_available("zarr"):
+        from modelrunner.storage.backend.zarr import is_zarr2
+
         exts.append("zip")
         if incl_folder:
             exts.extend(["", "zarr"])
-        if module_available("sqlite3"):
+        if is_zarr2 and module_available("sqlite3"):
             exts.append("sqldb")
 
     if dot: