allow using zarr2 for python 3.10, fix save_zarr code and test

Author: pre-commit-ci[bot]

pyproject.toml

@@ -59,7 +59,7 @@ dependencies = [
   "ruamel-yaml",
   "semantic-version",
   "tqdm",
-  "zarr>=3.1.3",
+  "zarr",                             # will use zarr3 for python >=3.11
 ]
 
 optional-dependencies.all = [

src/anemoi/datasets/data/misc.py

@@ -664,7 +664,7 @@ def initialize_zarr_store(root: Any, big_dataset: "Dataset") -> None:
     # Create or append to "dates" dataset.
     if "dates" not in root:
         full_length = len(big_dataset.dates)
-        root.create_dataset("dates", data=np.array([], dtype="datetime64[s]"), chunks=(full_length,))
+        root.create_dataset("dates", data=np.array([], dtype="datetime64[s]"), chunks=(full_length,), shape=(0,))
 
     if "data" not in root:
         dims = (1, len(big_dataset.variables), ensembles, big_dataset.shape[-1])
@@ -681,12 +681,15 @@ def initialize_zarr_store(root: Any, big_dataset: "Dataset") -> None:
                 k,
                 data=v,
                 compressor=None,
+                shape=v.shape,
             )
 
     # Create spatial coordinate datasets if missing.
     if "latitudes" not in root or "longitudes" not in root:
-        root.create_dataset("latitudes", data=big_dataset.latitudes, compressor=None)
-        root.create_dataset("longitudes", data=big_dataset.longitudes, compressor=None)
+        root.create_dataset("latitudes", data=big_dataset.latitudes, compressor=None, shape=big_dataset.latitudes.shape)
+        root.create_dataset(
+            "longitudes", data=big_dataset.longitudes, compressor=None, shape=big_dataset.longitudes.shape
+        )
     for k, v in big_dataset.metadata().items():
         if k not in root.attrs:
             root.attrs[k] = v

tests/create/utils/compare.py

@@ -161,7 +161,7 @@ def compare_dot_zattrs(a: dict, b: dict, path: str, errors: list) -> None:
                     "description",
                     "config_path",
                     "total_size",
-                    "total_number_of_files", # these are expected to differ between zarr 2 and zarr 3
+                    "total_number_of_files",  # these are expected to differ between zarr 2 and zarr 3
                 ]:
                     if type(a[k]) is not type(b[k]):
                         errors.append(f"❌ {path}.{k} : type differs {type(a[k])} != {type(b[k])}")

tests/test_data.py

@@ -42,6 +42,8 @@
 
 VALUES = 10
 
+true_zarr_open = zarr.open
+
 
 def mockup_open_zarr(func: Callable) -> Callable:
     """Decorator to mock the open_zarr function.
@@ -237,6 +239,8 @@ def zarr_from_str(name: str, mode: str) -> zarr.Group:
         Zarr dataset.
     """
     # Format: test-2021-2021-6h-o96-abcd-0
+    if "/" in name:
+        return true_zarr_open(name)
 
     args = dict(
         test="test",