diff --git a/virtualizarr/manifests/array.py b/virtualizarr/manifests/array.py index 393ddf2f..511e2972 100644 --- a/virtualizarr/manifests/array.py +++ b/virtualizarr/manifests/array.py @@ -1,9 +1,9 @@ import warnings -from typing import Any, Callable, Union, cast +from typing import TYPE_CHECKING, Any, Callable, Union, cast import numpy as np import xarray as xr -from zarr.core.metadata.v3 import ArrayV3Metadata, RegularChunkGrid +from zarr.core.metadata.v3 import ArrayV3Metadata import virtualizarr.manifests.utils as utils from virtualizarr.manifests.array_api import ( @@ -14,6 +14,20 @@ from virtualizarr.manifests.manifest import ChunkManifest from virtualizarr.manifests.utils import ChunkKeySeparator +# Type-check against the min-deps name (RegularChunkGrid) while keeping a +# runtime try/except so both zarr-python <=3.1.6 and >3.1.6 work. Version +# sniffing is unreliable under hatch-vcs when installed from a git source +# without fetched tags. +if TYPE_CHECKING: + from zarr.core.metadata.v3 import RegularChunkGrid as RegularChunkGridMetadata +else: + try: + from zarr.core.metadata.v3 import RegularChunkGridMetadata # zarr-python>3.1.6 + except ImportError: + from zarr.core.metadata.v3 import ( + RegularChunkGrid as RegularChunkGridMetadata, # zarr-python<=3.1.6 + ) + class ManifestArray: """ @@ -50,7 +64,7 @@ def __init__( # try unpacking the dict _metadata = ArrayV3Metadata(**metadata) - if not isinstance(_metadata.chunk_grid, RegularChunkGrid): + if not isinstance(_metadata.chunk_grid, RegularChunkGridMetadata): raise NotImplementedError( f"Only RegularChunkGrid is currently supported for chunk size, but got type {type(_metadata.chunk_grid)}" ) diff --git a/virtualizarr/parsers/zarr.py b/virtualizarr/parsers/zarr.py index b09ba81d..b753a191 100644 --- a/virtualizarr/parsers/zarr.py +++ b/virtualizarr/parsers/zarr.py @@ -6,14 +6,13 @@ from collections.abc import Coroutine, Iterable from enum import Enum from pathlib import Path -from typing import Any, TypeVar, cast +from typing import TYPE_CHECKING, Any, TypeVar, cast import numpy as np import obstore import zarr from obspec_utils.registry import ObjectStoreRegistry from zarr.api.asynchronous import open_group as open_group_async -from zarr.core.chunk_grids import RegularChunkGrid from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata from zarr.storage import ObjectStore @@ -29,6 +28,20 @@ from virtualizarr.manifests.utils import ChunkKeySeparator from virtualizarr.utils import determine_chunk_grid_shape +# Type-check against the min-deps name (RegularChunkGrid) while keeping a +# runtime try/except so both zarr-python <=3.1.6 and >3.1.6 work. Version +# sniffing is unreliable under hatch-vcs when installed from a git source +# without fetched tags. +if TYPE_CHECKING: + from zarr.core.metadata.v3 import RegularChunkGrid as RegularChunkGridMetadata +else: + try: + from zarr.core.metadata.v3 import RegularChunkGridMetadata # zarr-python>3.1.6 + except ImportError: + from zarr.core.metadata.v3 import ( + RegularChunkGrid as RegularChunkGridMetadata, # zarr-python<=3.1.6 + ) + # obstore doesn't export a public base type for stores, so we use Any for now. ObstoreStore = Any @@ -267,7 +280,7 @@ async def construct_manifest_array( """Construct a ManifestArray from a zarr array.""" array_v3_metadata = metadata_as_v3(zarr_array.metadata) - if not isinstance(array_v3_metadata.chunk_grid, RegularChunkGrid): + if not isinstance(array_v3_metadata.chunk_grid, RegularChunkGridMetadata): raise NotImplementedError( f"Only RegularChunkGrid is supported, but array {zarr_array.path} " f"uses {type(array_v3_metadata.chunk_grid).__name__}." @@ -387,7 +400,7 @@ async def build_chunk_manifest( # chunk shape, which lives inside the ShardingCodec config). So this grid describes # the number of shard files on disk, which is exactly what we want for the manifest. chunk_grid_shape = determine_chunk_grid_shape( - metadata.shape, cast(RegularChunkGrid, metadata.chunk_grid).chunk_shape + metadata.shape, cast(RegularChunkGridMetadata, metadata.chunk_grid).chunk_shape ) total_size = math.prod(chunk_grid_shape)