Skip to content

Commit 547836b

Browse files
d-v-bclaude
andcommitted
fix: widen is_regular_* annotations and cast ChunksTuple in create_chunk_grid_metadata
Accept ndarray[int64] in is_regular_1d/is_regular_nd alongside Sequence[int]. Cast only the first element per axis on the regular path so D ints are allocated rather than N*D. Materialize fully on the rectilinear path because _validate_chunk_shapes checks isinstance(dim_spec, int) which rejects np.int64. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 62ce735 commit 547836b

1 file changed

Lines changed: 12 additions & 4 deletions

File tree

  • src/zarr/core/metadata

src/zarr/core/metadata/v3.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@
3939
if TYPE_CHECKING:
4040
from typing import Self
4141

42+
import numpy as np
43+
4244
from zarr.core.buffer import Buffer, BufferPrototype
4345
from zarr.core.chunk_grids import ChunksTuple
4446
from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar
@@ -371,7 +373,9 @@ def from_dict(cls, data: RectilinearChunkGridMetadataJSON) -> Self: # type: ign
371373
ChunkGridMetadata = RegularChunkGridMetadata | RectilinearChunkGridMetadata
372374

373375

374-
def is_regular_1d(dim_chunks: Sequence[int]) -> bool:
376+
def is_regular_1d(
377+
dim_chunks: Sequence[int] | np.ndarray[tuple[int], np.dtype[np.int64]],
378+
) -> bool:
375379
"""Check if a single dimension's chunk sizes represent a regular grid.
376380
377381
A regular dimension has either all chunks the same size, or all
@@ -388,7 +392,9 @@ def is_regular_1d(dim_chunks: Sequence[int]) -> bool:
388392
return dim_chunks[-1] <= first
389393

390394

391-
def is_regular_nd(chunks: Iterable[Sequence[int]]) -> bool:
395+
def is_regular_nd(
396+
chunks: Iterable[Sequence[int] | np.ndarray[tuple[int], np.dtype[np.int64]]],
397+
) -> bool:
392398
"""Check if an N-dimensional chunk specification represents a regular grid."""
393399
return all(is_regular_1d(d) for d in chunks)
394400

@@ -414,10 +420,12 @@ def create_chunk_grid_metadata(
414420
if is_regular_nd(chunks):
415421
# If we know the chunks specification is regular, then we can take the first
416422
# chunk size for each dimension as the chunk shape.
417-
chunk_shape = tuple(dim_chunks[0] for dim_chunks in chunks)
423+
chunk_shape = tuple(int(dim_chunks[0]) for dim_chunks in chunks)
418424
return RegularChunkGridMetadata(chunk_shape=chunk_shape)
419425
else:
420-
return RectilinearChunkGridMetadata(chunk_shapes=chunks)
426+
return RectilinearChunkGridMetadata(
427+
chunk_shapes=tuple(tuple(int(x) for x in d) for d in chunks)
428+
)
421429

422430

423431
def parse_chunk_grid(

0 commit comments

Comments
 (0)