Skip to content

Commit

Permalink
format
Browse files Browse the repository at this point in the history
  • Loading branch information
brokkoli71 committed Oct 22, 2024
1 parent f063515 commit d8094ba
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 9 deletions.
9 changes: 5 additions & 4 deletions src/zarr/core/indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -1344,15 +1344,16 @@ def decode_morton(z: int, chunk_shape: ChunkCoords) -> ChunkCoords:

def morton_order_iter(chunk_shape: ChunkCoords) -> Iterator[ChunkCoords]:
i = 0
order = []
while len(order)<product(chunk_shape):
order: list[ChunkCoords] = []
while len(order) < product(chunk_shape):
m = decode_morton(i, chunk_shape)
if m not in order and all(x < y for x,y in zip(m, chunk_shape, strict=False)):
if m not in order and all(x < y for x, y in zip(m, chunk_shape, strict=False)):
order.append(m)
i+=1
i += 1
for j in range(product(chunk_shape)):
yield order[j]


def c_order_iter(chunks_per_shard: ChunkCoords) -> Iterator[ChunkCoords]:
return itertools.product(*(range(x) for x in chunks_per_shard))

Expand Down
14 changes: 9 additions & 5 deletions tests/test_codecs/test_sharding.py
Original file line number Diff line number Diff line change
Expand Up @@ -394,14 +394,18 @@ async def test_sharding_with_empty_inner_chunk(
data_read = await a.getitem(...)
assert np.array_equal(data_read, data)


@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
@pytest.mark.parametrize(
"index_location", [ShardingCodecIndexLocation.start, ShardingCodecIndexLocation.end],
"index_location",
[ShardingCodecIndexLocation.start, ShardingCodecIndexLocation.end],
)
@pytest.mark.parametrize("chunks_per_shard", [(5,2), (2,5), (5,5)])
async def test_sharding_with_chunks_per_shard(store: Store, index_location: ShardingCodecIndexLocation, chunks_per_shard: tuple) -> None:
chunk_shape = (2,1)
shape = [x*y for x,y in zip(chunks_per_shard, chunk_shape, strict=False)]
@pytest.mark.parametrize("chunks_per_shard", [(5, 2), (2, 5), (5, 5)])
async def test_sharding_with_chunks_per_shard(
store: Store, index_location: ShardingCodecIndexLocation, chunks_per_shard: tuple[int]
) -> None:
chunk_shape = (2, 1)
shape = [x * y for x, y in zip(chunks_per_shard, chunk_shape, strict=False)]
data = np.ones(np.prod(shape), dtype="int32").reshape(shape)
fill_value = 42

Expand Down

0 comments on commit d8094ba

Please sign in to comment.