|
2 | 2 | import math |
3 | 3 | import pickle |
4 | 4 | import warnings |
5 | | -from typing import Any, cast, get_args |
| 5 | +from typing import TYPE_CHECKING, Any, cast, get_args |
6 | 6 | from unittest.mock import AsyncMock |
7 | 7 |
|
8 | 8 | import numpy as np |
|
40 | 40 | from ..conftest import ArrayRequest |
41 | 41 | from .test_codecs import _AsyncArrayProxy, order_from_dim |
42 | 42 |
|
| 43 | +if TYPE_CHECKING: |
| 44 | + from zarr.core.common import JSON |
| 45 | + |
43 | 46 |
|
44 | 47 | def _reads_are_sync(store_mock: AsyncMock) -> bool: |
45 | 48 | """True when the partial-shard read for this store+pipeline goes through the |
@@ -765,15 +768,13 @@ def test_structured_dtype_fill_value() -> None: |
765 | 768 | [ |
766 | 769 | ShardingCodec(chunk_shape=(8, 8)), |
767 | 770 | ShardingCodec(chunk_shape=(8, 8), subchunk_write_order="lexicographic"), |
768 | | - ShardingCodec(chunk_shape=(0,)), |
769 | 771 | ], |
770 | | - ids=["default", "lexicographic", "zero-chunk-size"], |
| 772 | + ids=["default", "lexicographic"], |
771 | 773 | ) |
772 | 774 | def test_pickle(codec: ShardingCodec) -> None: |
773 | 775 | """ShardingCodec round-trips through pickle, including the non-serialized |
774 | 776 | ``subchunk_write_order`` (which ``to_dict`` omits and which must not silently |
775 | | - revert to the ``morton`` default), and an inner chunk size of 0, which the |
776 | | - constructor accepts.""" |
| 777 | + revert to the ``morton`` default).""" |
777 | 778 | restored = pickle.loads(pickle.dumps(codec)) |
778 | 779 | assert restored == codec |
779 | 780 | assert restored.chunk_shape == codec.chunk_shape |
@@ -1395,3 +1396,51 @@ def test_sharding_orthogonal_set_multiple_array_dims( |
1395 | 1396 | expected[ix] = value.reshape(expected[ix].shape) |
1396 | 1397 | assert np.array_equal(a[:], expected) |
1397 | 1398 | assert np.array_equal(a.oindex[selection], value) |
| 1399 | + |
| 1400 | + |
| 1401 | +@pytest.mark.parametrize( |
| 1402 | + ("chunk_shape", "expected"), |
| 1403 | + [ |
| 1404 | + ((4,), (4,)), |
| 1405 | + ([4, 2], (4, 2)), |
| 1406 | + (4, (4,)), |
| 1407 | + (np.int64(4), (4,)), |
| 1408 | + ((np.int64(4),), (4,)), |
| 1409 | + (np.array([4]), (4,)), |
| 1410 | + ((True,), (1,)), |
| 1411 | + (range(4, 5), (4,)), |
| 1412 | + ((), ()), |
| 1413 | + ], |
| 1414 | + ids=repr, |
| 1415 | +) |
| 1416 | +def test_sharding_codec_chunk_shape(chunk_shape: Any, expected: tuple[int, ...]) -> None: |
| 1417 | + """The inner chunk shape of a shard is parsed by the chunk normalizer's integer test as |
| 1418 | + a regular chunk shape: one integer of at least 1 per dimension, written back as |
| 1419 | + given.""" |
| 1420 | + codec = ShardingCodec(chunk_shape=chunk_shape) |
| 1421 | + assert codec.chunk_shape == expected |
| 1422 | + assert all(type(size) is int for size in codec.chunk_shape) |
| 1423 | + configuration = cast("dict[str, JSON]", codec.to_dict()["configuration"]) |
| 1424 | + assert configuration["chunk_shape"] == expected |
| 1425 | + assert ShardingCodec.from_dict(codec.to_dict()).chunk_shape == expected |
| 1426 | + |
| 1427 | + |
| 1428 | +@pytest.mark.parametrize("chunk_shape", [(0,), (False,), (-1,), 0, [8, 0]], ids=repr) |
| 1429 | +def test_sharding_codec_chunk_shape_rejects_non_positive(chunk_shape: Any) -> None: |
| 1430 | + """An inner chunk size below 1 is rejected, including `-1` and `False`, which have no |
| 1431 | + axis length to be resolved against.""" |
| 1432 | + with pytest.raises(ValueError, match="All chunk sizes must be positive"): |
| 1433 | + ShardingCodec(chunk_shape=chunk_shape) |
| 1434 | + |
| 1435 | + |
| 1436 | +@pytest.mark.parametrize("chunk_shape", [[[4, 4]], [8, [4, 4]], (4.0,), "4"], ids=repr) |
| 1437 | +def test_sharding_codec_chunk_shape_rejects_non_integer(chunk_shape: Any) -> None: |
| 1438 | + """An inner chunk shape is regular: an explicit list of chunk edges, which declares a |
| 1439 | + rectilinear dimension, is rejected like any other non-integer chunk size.""" |
| 1440 | + with pytest.raises(TypeError, match="Each chunk size of a regular chunk shape must be"): |
| 1441 | + ShardingCodec(chunk_shape=chunk_shape) |
| 1442 | + |
| 1443 | + |
| 1444 | +def test_sharding_codec_chunk_shape_rejects_non_iterable() -> None: |
| 1445 | + with pytest.raises(TypeError, match="must be an integer or an iterable of integers"): |
| 1446 | + ShardingCodec(chunk_shape=None) # type: ignore[arg-type] |
0 commit comments