From 2efb368f12808d782e70ab13943f4c33c56f3eac Mon Sep 17 00:00:00 2001
From: Deepak Cherian <deepak@cherian.net>
Date: Sun, 26 Jan 2025 19:48:30 -0700
Subject: [PATCH 1/9] Fix some typing

---
 xarray/backends/zarr.py | 6 ++++--
 xarray/core/types.py    | 2 +-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py
index 383c385e1d5..20364c2d980 100644
--- a/xarray/backends/zarr.py
+++ b/xarray/backends/zarr.py
@@ -768,7 +768,7 @@ def __init__(
             self._members = self._fetch_members()
 
     @property
-    def members(self) -> dict[str, ZarrArray]:
+    def members(self) -> dict[str, ZarrArray | ZarrGroup]:
         """
         Model the arrays and groups contained in self.zarr_group as a dict. If `self._cache_members`
         is true, the dict is cached. Otherwise, it is retrieved from storage.
@@ -778,7 +778,7 @@ def members(self) -> dict[str, ZarrArray]:
         else:
             return self._members
 
-    def _fetch_members(self) -> dict[str, ZarrArray]:
+    def _fetch_members(self) -> dict[str, ZarrArray | ZarrGroup]:
         """
         Get the arrays and groups defined in the zarr group modelled by this Store
         """
@@ -1066,6 +1066,8 @@ def _open_existing_array(self, *, name) -> ZarrArray:
         else:
             zarr_array = self.zarr_group[name]
 
+        if TYPE_CHECKING:
+            assert isinstance(zarr_array, ZarrArray)
         return zarr_array
 
     def _create_new_array(
diff --git a/xarray/core/types.py b/xarray/core/types.py
index d4224dcead9..4232433cf19 100644
--- a/xarray/core/types.py
+++ b/xarray/core/types.py
@@ -71,7 +71,7 @@
     try:
         from zarr import Array as ZarrArray
     except ImportError:
-        ZarrArray = np.ndarray
+        ZarrArray = np.ndarray  # type: ignore[misc, assignment, unused-ignore]
 
     # Anything that can be coerced to a shape tuple
     _ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]]

From 8de93cb871469cc12b8dc0a4a8a5e6f5a31439c5 Mon Sep 17 00:00:00 2001
From: Deepak Cherian <deepak@cherian.net>
Date: Mon, 27 Jan 2025 08:59:04 -0700
Subject: [PATCH 2/9] Use StoreLike

---
 xarray/backends/api.py     | 8 ++++----
 xarray/core/datatree_io.py | 6 +++---
 xarray/core/types.py       | 2 ++
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/xarray/backends/api.py b/xarray/backends/api.py
index 3211b9efbae..5a145b37686 100644
--- a/xarray/backends/api.py
+++ b/xarray/backends/api.py
@@ -45,7 +45,7 @@
 from xarray.core.datatree import DataTree
 from xarray.core.indexes import Index
 from xarray.core.treenode import group_subtrees
-from xarray.core.types import NetcdfWriteModes, ZarrWriteModes
+from xarray.core.types import NetcdfWriteModes, ZarrStoreLike, ZarrWriteModes
 from xarray.core.utils import is_remote_uri
 from xarray.namedarray.daskmanager import DaskManager
 from xarray.namedarray.parallelcompat import guess_chunkmanager
@@ -2100,7 +2100,7 @@ def save_mfdataset(
 @overload
 def to_zarr(
     dataset: Dataset,
-    store: MutableMapping | str | os.PathLike[str] | None = None,
+    store: ZarrStoreLike | None = None,
     chunk_store: MutableMapping | str | os.PathLike | None = None,
     mode: ZarrWriteModes | None = None,
     synchronizer=None,
@@ -2123,7 +2123,7 @@ def to_zarr(
 @overload
 def to_zarr(
     dataset: Dataset,
-    store: MutableMapping | str | os.PathLike[str] | None = None,
+    store: ZarrStoreLike | None = None,
     chunk_store: MutableMapping | str | os.PathLike | None = None,
     mode: ZarrWriteModes | None = None,
     synchronizer=None,
@@ -2144,7 +2144,7 @@ def to_zarr(
 
 def to_zarr(
     dataset: Dataset,
-    store: MutableMapping | str | os.PathLike[str] | None = None,
+    store: ZarrStoreLike | None = None,
     chunk_store: MutableMapping | str | os.PathLike | None = None,
     mode: ZarrWriteModes | None = None,
     synchronizer=None,
diff --git a/xarray/core/datatree_io.py b/xarray/core/datatree_io.py
index 3d0daa26b90..f28c5d8319b 100644
--- a/xarray/core/datatree_io.py
+++ b/xarray/core/datatree_io.py
@@ -1,11 +1,11 @@
 from __future__ import annotations
 
-from collections.abc import Mapping, MutableMapping
+from collections.abc import Mapping
 from os import PathLike
 from typing import Any, Literal, get_args
 
 from xarray.core.datatree import DataTree
-from xarray.core.types import NetcdfWriteModes, ZarrWriteModes
+from xarray.core.types import NetcdfWriteModes, ZarrStoreLike, ZarrWriteModes
 
 T_DataTreeNetcdfEngine = Literal["netcdf4", "h5netcdf"]
 T_DataTreeNetcdfTypes = Literal["NETCDF4"]
@@ -78,7 +78,7 @@ def _datatree_to_netcdf(
 
 def _datatree_to_zarr(
     dt: DataTree,
-    store: MutableMapping | str | PathLike[str],
+    store: ZarrStoreLike,
     mode: ZarrWriteModes = "w-",
     encoding: Mapping[str, Any] | None = None,
     consolidated: bool = True,
diff --git a/xarray/core/types.py b/xarray/core/types.py
index 4232433cf19..0372043e573 100644
--- a/xarray/core/types.py
+++ b/xarray/core/types.py
@@ -70,8 +70,10 @@
 
     try:
         from zarr import Array as ZarrArray
+        from zarr.storage import StoreLike as ZarrStoreLike
     except ImportError:
         ZarrArray = np.ndarray  # type: ignore[misc, assignment, unused-ignore]
+        ZarrStoreLike = Any  # type: ignore[misc]
 
     # Anything that can be coerced to a shape tuple
     _ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]]

From e838999b4ae6d387833933a152ef9f56c8403dde Mon Sep 17 00:00:00 2001
From: Deepak Cherian <deepak@cherian.net>
Date: Mon, 27 Jan 2025 09:01:03 -0700
Subject: [PATCH 3/9] tweak

---
 xarray/backends/zarr.py | 6 ++----
 xarray/core/types.py    | 2 +-
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py
index 20364c2d980..8ea1e5e99c9 100644
--- a/xarray/backends/zarr.py
+++ b/xarray/backends/zarr.py
@@ -5,7 +5,7 @@
 import os
 import struct
 from collections.abc import Hashable, Iterable, Mapping
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
 import pandas as pd
@@ -1066,9 +1066,7 @@ def _open_existing_array(self, *, name) -> ZarrArray:
         else:
             zarr_array = self.zarr_group[name]
 
-        if TYPE_CHECKING:
-            assert isinstance(zarr_array, ZarrArray)
-        return zarr_array
+        return cast(ZarrArray, zarr_array)
 
     def _create_new_array(
         self, *, name, shape, dtype, fill_value, encoding, attrs
diff --git a/xarray/core/types.py b/xarray/core/types.py
index 0372043e573..0201e5322d6 100644
--- a/xarray/core/types.py
+++ b/xarray/core/types.py
@@ -73,7 +73,7 @@
         from zarr.storage import StoreLike as ZarrStoreLike
     except ImportError:
         ZarrArray = np.ndarray  # type: ignore[misc, assignment, unused-ignore]
-        ZarrStoreLike = Any  # type: ignore[misc]
+        ZarrStoreLike = Any  # type: ignore[misc, assignment, unused-ignore]
 
     # Anything that can be coerced to a shape tuple
     _ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]]

From 03798d036eb366fe38969e0e235d4677bc10a725 Mon Sep 17 00:00:00 2001
From: Deepak Cherian <deepak@cherian.net>
Date: Mon, 27 Jan 2025 09:13:49 -0700
Subject: [PATCH 4/9] fix

---
 xarray/core/types.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/xarray/core/types.py b/xarray/core/types.py
index 0201e5322d6..82d934ebf75 100644
--- a/xarray/core/types.py
+++ b/xarray/core/types.py
@@ -70,10 +70,8 @@
 
     try:
         from zarr import Array as ZarrArray
-        from zarr.storage import StoreLike as ZarrStoreLike
     except ImportError:
         ZarrArray = np.ndarray  # type: ignore[misc, assignment, unused-ignore]
-        ZarrStoreLike = Any  # type: ignore[misc, assignment, unused-ignore]
 
     # Anything that can be coerced to a shape tuple
     _ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]]
@@ -117,6 +115,13 @@
 )
 
 
+try:
+    # this is V3 only
+    from zarr.storage import StoreLike as ZarrStoreLike
+except ImportError:
+    ZarrStoreLike = Any  # type: ignore[misc, assignment, unused-ignore]
+
+
 class Alignable(Protocol):
     """Represents any Xarray type that supports alignment.
 

From baccc1cf2393049fa289a723e13be231ab3608cd Mon Sep 17 00:00:00 2001
From: Deepak Cherian <deepak@cherian.net>
Date: Mon, 27 Jan 2025 09:23:59 -0700
Subject: [PATCH 5/9] one ignore

---
 xarray/tests/test_backends_datatree.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/xarray/tests/test_backends_datatree.py b/xarray/tests/test_backends_datatree.py
index 608f9645ce8..ff84cf4c0ea 100644
--- a/xarray/tests/test_backends_datatree.py
+++ b/xarray/tests/test_backends_datatree.py
@@ -414,7 +414,7 @@ def test_to_zarr_zip_store(self, tmpdir, simple_datatree):
         store = ZipStore(filepath)
         original_dt.to_zarr(store)
 
-        with open_datatree(store, engine="zarr") as roundtrip_dt:
+        with open_datatree(store, engine="zarr") as roundtrip_dt:  # type: ignore[arg-type]
             assert_equal(original_dt, roundtrip_dt)
 
     def test_to_zarr_not_consolidated(self, tmpdir, simple_datatree):

From 59383091e75a804595cd888990c0ad982f369379 Mon Sep 17 00:00:00 2001
From: Deepak Cherian <deepak@cherian.net>
Date: Mon, 27 Jan 2025 09:27:54 -0700
Subject: [PATCH 6/9] one more ignore

---
 xarray/tests/test_backends_datatree.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/xarray/tests/test_backends_datatree.py b/xarray/tests/test_backends_datatree.py
index ff84cf4c0ea..9cbe3f2247a 100644
--- a/xarray/tests/test_backends_datatree.py
+++ b/xarray/tests/test_backends_datatree.py
@@ -414,7 +414,7 @@ def test_to_zarr_zip_store(self, tmpdir, simple_datatree):
         store = ZipStore(filepath)
         original_dt.to_zarr(store)
 
-        with open_datatree(store, engine="zarr") as roundtrip_dt:  # type: ignore[arg-type]
+        with open_datatree(store, engine="zarr") as roundtrip_dt:  # type: ignore[arg-type, unused-ignore]
             assert_equal(original_dt, roundtrip_dt)
 
     def test_to_zarr_not_consolidated(self, tmpdir, simple_datatree):

From 1e434c8d2e56859741647297e054fc014458afaa Mon Sep 17 00:00:00 2001
From: Deepak Cherian <deepak@cherian.net>
Date: Mon, 27 Jan 2025 09:52:37 -0700
Subject: [PATCH 7/9] try again

---
 xarray/backends/zarr.py | 6 ++----
 xarray/core/types.py    | 2 ++
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py
index 8ea1e5e99c9..16e141c62c3 100644
--- a/xarray/backends/zarr.py
+++ b/xarray/backends/zarr.py
@@ -38,13 +38,10 @@
 from xarray.namedarray.utils import module_available
 
 if TYPE_CHECKING:
-    from zarr import Array as ZarrArray
-    from zarr import Group as ZarrGroup
-
     from xarray.backends.common import AbstractDataStore
     from xarray.core.dataset import Dataset
     from xarray.core.datatree import DataTree
-    from xarray.core.types import ReadBuffer
+    from xarray.core.types import ReadBuffer, ZarrArray, ZarrGroup
 
 
 def _get_mappers(*, storage_options, store, chunk_store):
@@ -1035,6 +1032,7 @@ def sync(self):
 
     def _open_existing_array(self, *, name) -> ZarrArray:
         import zarr
+        from zarr import Array as ZarrArray
 
         # TODO: if mode="a", consider overriding the existing variable
         # metadata. This would need some case work properly with region
diff --git a/xarray/core/types.py b/xarray/core/types.py
index 82d934ebf75..0199e8dd502 100644
--- a/xarray/core/types.py
+++ b/xarray/core/types.py
@@ -70,8 +70,10 @@
 
     try:
         from zarr import Array as ZarrArray
+        from zarr import Group as ZarrGroup
     except ImportError:
         ZarrArray = np.ndarray  # type: ignore[misc, assignment, unused-ignore]
+        ZarrGroup = Any  # type: ignore[misc, assignment, unused-ignore]
 
     # Anything that can be coerced to a shape tuple
     _ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]]

From c96c846115fee7315b5a1c483108c8461c3d4d52 Mon Sep 17 00:00:00 2001
From: Deepak Cherian <deepak@cherian.net>
Date: Mon, 27 Jan 2025 09:57:26 -0700
Subject: [PATCH 8/9] fix _zarr_v3

---
 xarray/backends/zarr.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py
index 16e141c62c3..a584f78e63a 100644
--- a/xarray/backends/zarr.py
+++ b/xarray/backends/zarr.py
@@ -105,8 +105,7 @@ def _choose_default_mode(
 
 
 def _zarr_v3() -> bool:
-    # TODO: switch to "3" once Zarr V3 is released
-    return module_available("zarr", minversion="2.99")
+    return module_available("zarr", minversion="3")
 
 
 # need some special secret attributes to tell us the dimensions

From b2b6530e76e7dafd3da2f80f2927d9aab715940c Mon Sep 17 00:00:00 2001
From: Deepak Cherian <deepak@cherian.net>
Date: Mon, 27 Jan 2025 10:10:04 -0700
Subject: [PATCH 9/9] try again

---
 xarray/backends/api.py     |  3 ++-
 xarray/core/datatree_io.py |  7 +++++--
 xarray/core/types.py       | 12 +++++-------
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/xarray/backends/api.py b/xarray/backends/api.py
index 5a145b37686..ef97bfb8205 100644
--- a/xarray/backends/api.py
+++ b/xarray/backends/api.py
@@ -45,7 +45,7 @@
 from xarray.core.datatree import DataTree
 from xarray.core.indexes import Index
 from xarray.core.treenode import group_subtrees
-from xarray.core.types import NetcdfWriteModes, ZarrStoreLike, ZarrWriteModes
+from xarray.core.types import NetcdfWriteModes, ZarrWriteModes
 from xarray.core.utils import is_remote_uri
 from xarray.namedarray.daskmanager import DaskManager
 from xarray.namedarray.parallelcompat import guess_chunkmanager
@@ -64,6 +64,7 @@
         NestedSequence,
         ReadBuffer,
         T_Chunks,
+        ZarrStoreLike,
     )
 
     T_NetcdfEngine = Literal["netcdf4", "scipy", "h5netcdf"]
diff --git a/xarray/core/datatree_io.py b/xarray/core/datatree_io.py
index f28c5d8319b..3a57ebc7d2d 100644
--- a/xarray/core/datatree_io.py
+++ b/xarray/core/datatree_io.py
@@ -2,14 +2,17 @@
 
 from collections.abc import Mapping
 from os import PathLike
-from typing import Any, Literal, get_args
+from typing import TYPE_CHECKING, Any, Literal, get_args
 
 from xarray.core.datatree import DataTree
-from xarray.core.types import NetcdfWriteModes, ZarrStoreLike, ZarrWriteModes
+from xarray.core.types import NetcdfWriteModes, ZarrWriteModes
 
 T_DataTreeNetcdfEngine = Literal["netcdf4", "h5netcdf"]
 T_DataTreeNetcdfTypes = Literal["NETCDF4"]
 
+if TYPE_CHECKING:
+    from xarray.core.types import ZarrStoreLike
+
 
 def _datatree_to_netcdf(
     dt: DataTree,
diff --git a/xarray/core/types.py b/xarray/core/types.py
index 0199e8dd502..186738ed718 100644
--- a/xarray/core/types.py
+++ b/xarray/core/types.py
@@ -74,6 +74,11 @@
     except ImportError:
         ZarrArray = np.ndarray  # type: ignore[misc, assignment, unused-ignore]
         ZarrGroup = Any  # type: ignore[misc, assignment, unused-ignore]
+    try:
+        # this is V3 only
+        from zarr.storage import StoreLike as ZarrStoreLike
+    except ImportError:
+        ZarrStoreLike = Any  # type: ignore[misc, assignment, unused-ignore]
 
     # Anything that can be coerced to a shape tuple
     _ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]]
@@ -117,13 +122,6 @@
 )
 
 
-try:
-    # this is V3 only
-    from zarr.storage import StoreLike as ZarrStoreLike
-except ImportError:
-    ZarrStoreLike = Any  # type: ignore[misc, assignment, unused-ignore]
-
-
 class Alignable(Protocol):
     """Represents any Xarray type that supports alignment.