add basic ChunkGrid tests
MODIFY
pangeo_forge_recipes/chunk_grid.py
pangeo_forge_recipes/chunk_grid.py
@@ -3,7 +3,7 @@ Abstract representation of ND chunked arrays
"""
from dataclasses import dataclass
-from typing import Tuple
+from typing import Dict, FrozenSet, Tuple
import numpy as np
@@ -13,30 +13,25 @@ import numpy as np
ChunkIndex = Tuple[int, ...]
-@dataclass
class ChunkGrid:
- shape: Tuple[int, ...]
+ def __init__(self, chunks):
+ self._chunk_axes = {name: ChunkAxis(axis_chunks) for name, axis_chunks in chunks.items()}
- # we will always be dealing with named dims
- dims: Tuple[str, ...]
+ @property
+ def dims(self) -> FrozenSet[str]:
+ return frozenset(self._chunk_axes)
- # chunks are stored explicitly (don't assume uniform length)
- chunks: Tuple[Tuple[int, ...], ...]
+ @property
+ def shape(self) -> Dict[str, int]:
+ return {name: len(ca) for name, ca in self._chunk_axes.items()}
@property
- def ndim(self):
- return len(self.shape)
+ def nchunks(self) -> Dict[str, int]:
+ return {name: ca.nchunks for name, ca in self._chunk_axes.items()}
- def __post_init__(self):
- if len(self.dims) != self.ndim: # pragma: no cover
- raise ValueError("dims must have the same len as shape")
- if len(self.chunks) != self.ndim: # pragma: no cover
- raise ValueError("chunks must have the same len as shape")
-
- self._chunk_axes = [
- ChunkAxis(dim_shape, dim_chunks)
- for dim_shape, dim_chunks in zip(self.shape, self.chunks)
- ]
+ @property
+ def ndim(self):
+ return len(self._chunk_axes)
def chunk_index_to_slice(self, chunk_index: ChunkIndex) -> Tuple[slice, ...]:
return tuple(

MODIFY
tests/test_chunk_grid.py
tests/test_chunk_grid.py
import pytest
-from pangeo_forge_recipes.chunk_grid import ChunkAxis
+from pangeo_forge_recipes.chunk_grid import ChunkAxis, ChunkGrid
def test_chunk_axis():
@@ -35,3 +35,11 @@ def test_chunk_axis():
assert ca.array_slice_to_chunk_slice(slice(2, 6)) == slice(1, 2)
assert ca.array_slice_to_chunk_slice(slice(2, 5)) == slice(1, 2)
assert ca.array_slice_to_chunk_slice(slice(6, 7)) == slice(2, 3)
+
+
+def test_chunk_grid():
+ cg = ChunkGrid({"x": (2, 4, 3), "time": (7, 8)})
+ assert cg.dims == {"x", "time"}
+ assert cg.shape == {"x": 9, "time": 15}
+ assert cg.nchunks == {"x": 3, "time": 2}
+ assert cg.ndim == 2

add basic ChunkGrid tests