fix typing errors

after zarr-developers/zarr-python#3304 we need
to be specific about if we have V2 or V3 arrays

Author: ianhi

icechunk-python/python/icechunk/dask.py

@@ -1,6 +1,6 @@
 import functools
 from collections.abc import Callable, Mapping
-from typing import Any, ParamSpec, TypeAlias, TypeVar
+from typing import TYPE_CHECKING, Any, ParamSpec, TypeAlias, TypeVar
 
 import numpy as np
 from packaging.version import Version
@@ -13,6 +13,9 @@
 from icechunk.distributed import extract_session, merge_sessions
 from icechunk.session import ForkSession
 
+if TYPE_CHECKING:
+    from zarr.core.metadata import ArrayV3Metadata
+
 SimpleGraph: TypeAlias = Mapping[tuple[str, int], tuple[Any, ...]]
 
 
@@ -57,7 +60,7 @@ def _assert_correct_dask_version() -> None:
 def store_dask(
     *,
     sources: list[Array],
-    targets: list[zarr.Array],
+    targets: list[zarr.Array[ArrayV3Metadata]],
     regions: list[tuple[slice, ...]] | None = None,
     split_every: int | None = None,
     **store_kwargs: Any,

icechunk-python/python/icechunk/distributed.py

@@ -1,11 +1,14 @@
 # distributed utility functions
 from collections.abc import Generator, Iterable
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import zarr
 from icechunk import IcechunkStore
 from icechunk.session import ForkSession, Session
 
+if TYPE_CHECKING:
+    from zarr.core.metadata import ArrayV3Metadata
+
 __all__ = [
     "extract_session",
     "merge_sessions",
@@ -25,7 +28,7 @@ def _flatten(seq: Iterable[Any], container: type = list) -> Generator[Any, None,
 
 
 def extract_session(
-    zarray: zarr.Array, axis: Any = None, keepdims: Any = None
+    zarray: zarr.Array[ArrayV3Metadata], axis: Any = None, keepdims: Any = None
 ) -> Session:
     """
     Extract Icechunk Session from a zarr Array, useful for distributed array computing frameworks.

icechunk-python/python/icechunk/testing/strategies.py

@@ -1,16 +1,18 @@
 from collections.abc import Iterable
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import hypothesis.strategies as st
 
 import icechunk as ic
 import zarr
-from zarr.core.metadata import ArrayV3Metadata
+
+if TYPE_CHECKING:
+    from zarr.core.metadata import ArrayV3Metadata
 
 
 @st.composite
 def splitting_configs(
-    draw: st.DrawFn, *, arrays: Iterable[zarr.Array]
+    draw: st.DrawFn, *, arrays: Iterable[zarr.Array[ArrayV3Metadata]]
 ) -> ic.ManifestSplittingConfig:
     config_dict: dict[
         ic.ManifestSplitCondition,

icechunk-python/python/icechunk/xarray.py

@@ -1,7 +1,7 @@
 from collections.abc import Hashable, Mapping, MutableMapping
 from dataclasses import dataclass, field
 from importlib.util import find_spec
-from typing import Any, Literal, overload
+from typing import TYPE_CHECKING, Any, Literal, overload
 
 import numpy as np
 from packaging.version import Version
@@ -15,6 +15,9 @@
 from xarray.backends.common import ArrayWriter
 from xarray.backends.zarr import ZarrStore
 
+if TYPE_CHECKING:
+    from zarr.core.metadata import ArrayV3Metadata
+
 __all__ = ["to_icechunk"]
 
 Region = Mapping[str, slice | Literal["auto"]] | Literal["auto"] | None
@@ -57,7 +60,7 @@ def __init__(self) -> None:
         super().__init__()  # type: ignore[no-untyped-call]
 
         self.eager_sources: list[np.ndarray[Any, Any]] = []
-        self.eager_targets: list[zarr.Array] = []
+        self.eager_targets: list[zarr.Array[ArrayV3Metadata]] = []
         self.eager_regions: list[tuple[slice, ...]] = []
 
     def add(self, source: Any, target: Any, region: Any = None) -> Any: