diff --git a/icechunk-python/python/icechunk/dask.py b/icechunk-python/python/icechunk/dask.py index b4f24916c..ff8eac6c0 100644 --- a/icechunk-python/python/icechunk/dask.py +++ b/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, diff --git a/icechunk-python/python/icechunk/distributed.py b/icechunk-python/python/icechunk/distributed.py index c1a2608e8..91158408d 100644 --- a/icechunk-python/python/icechunk/distributed.py +++ b/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. diff --git a/icechunk-python/python/icechunk/testing/strategies.py b/icechunk-python/python/icechunk/testing/strategies.py index 435c374d2..c02b56fdd 100644 --- a/icechunk-python/python/icechunk/testing/strategies.py +++ b/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, diff --git a/icechunk-python/python/icechunk/xarray.py b/icechunk-python/python/icechunk/xarray.py index 6fd53e8ed..30a4be4a5 100644 --- a/icechunk-python/python/icechunk/xarray.py +++ b/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: