refactor: no enums

Author: ilan-gold

src/zarr/codecs/sharding.py

@@ -6,7 +6,7 @@
 from enum import Enum
 from functools import lru_cache
 from operator import itemgetter
-from typing import TYPE_CHECKING, Any, NamedTuple, cast
+from typing import TYPE_CHECKING, Any, Literal, NamedTuple, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -60,7 +60,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
-    from typing import Self
+    from typing import Final, Self
 
     from zarr.core.common import JSON
     from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType
@@ -79,27 +79,19 @@ class ShardingCodecIndexLocation(Enum):
     end = "end"
 
 
-class SubchunkWriteOrder(Enum):
-    """
-    Enum for the order of the chunks within a shard.
-
-    unordered is implemented via `random.shuffle` over the lexicographic order.
-    """
-
-    morton = "morton"
-    unordered = "unordered"
-    lexicographic = "lexicographic"
-    colexicographic = "colexicographic"
+SubchunkWriteOrder = Literal["morton", "unordered", "lexicographic", "colexicographic"]
+SUBCHUNK_WRITE_ORDER: Final[tuple[str, str, str, str]] = (
+    "morton",
+    "unordered",
+    "lexicographic",
+    "colexicographic",
+)
 
 
 def parse_index_location(data: object) -> ShardingCodecIndexLocation:
     return parse_enum(data, ShardingCodecIndexLocation)
 
 
-def parse_subchunk_write_order(data: object) -> SubchunkWriteOrder:
-    return parse_enum(data, SubchunkWriteOrder)
-
-
 @dataclass(frozen=True)
 class _ShardingByteGetter(ByteGetter):
     shard_dict: ShardMapping
@@ -324,7 +316,7 @@ class ShardingCodec(
     codecs: tuple[Codec, ...]
     index_codecs: tuple[Codec, ...]
     index_location: ShardingCodecIndexLocation = ShardingCodecIndexLocation.end
-    subchunk_write_order: SubchunkWriteOrder = SubchunkWriteOrder.morton
+    subchunk_write_order: SubchunkWriteOrder = "morton"
 
     def __init__(
         self,
@@ -333,19 +325,22 @@ def __init__(
         codecs: Iterable[Codec | dict[str, JSON]] = (BytesCodec(),),
         index_codecs: Iterable[Codec | dict[str, JSON]] = (BytesCodec(), Crc32cCodec()),
         index_location: ShardingCodecIndexLocation | str = ShardingCodecIndexLocation.end,
-        subchunk_write_order: SubchunkWriteOrder | str = SubchunkWriteOrder.morton,
+        subchunk_write_order: SubchunkWriteOrder = "morton",
     ) -> None:
         chunk_shape_parsed = parse_shapelike(chunk_shape)
         codecs_parsed = parse_codecs(codecs)
         index_codecs_parsed = parse_codecs(index_codecs)
         index_location_parsed = parse_index_location(index_location)
-        subchunk_write_order_parsed = parse_subchunk_write_order(subchunk_write_order)
+        if subchunk_write_order not in SUBCHUNK_WRITE_ORDER:
+            raise ValueError(
+                f"Unrecognized subchunk write order: {subchunk_write_order}. Only {SUBCHUNK_WRITE_ORDER} are allowed."
+            )
 
         object.__setattr__(self, "chunk_shape", chunk_shape_parsed)
         object.__setattr__(self, "codecs", codecs_parsed)
         object.__setattr__(self, "index_codecs", index_codecs_parsed)
         object.__setattr__(self, "index_location", index_location_parsed)
-        object.__setattr__(self, "subchunk_write_order", subchunk_write_order_parsed)
+        object.__setattr__(self, "subchunk_write_order", subchunk_write_order)
 
         # Use instance-local lru_cache to avoid memory leaks
 
@@ -547,21 +542,21 @@ async def _decode_partial_single(
 
     def _subchunk_order_iter(self, chunks_per_shard: tuple[int, ...]) -> Iterable[tuple[int, ...]]:
         match self.subchunk_write_order:
-            case SubchunkWriteOrder.morton:
+            case "morton":
                 subchunk_iter = morton_order_iter(chunks_per_shard)
-            case SubchunkWriteOrder.lexicographic:
+            case "lexicographic":
                 subchunk_iter = np.ndindex(chunks_per_shard)
-            case SubchunkWriteOrder.colexicographic:
+            case "colexicographic":
                 subchunk_iter = (c[::-1] for c in np.ndindex(chunks_per_shard[::-1]))
-            case SubchunkWriteOrder.unordered:
+            case "unordered":
                 subchunk_list = list(np.ndindex(chunks_per_shard))
                 random.shuffle(subchunk_list)
                 subchunk_iter = iter(subchunk_list)
         return subchunk_iter
 
     def _subchunk_order_vectorized(self, chunks_per_shard: tuple[int, ...]) -> npt.NDArray[np.intp]:
         match self.subchunk_write_order:
-            case SubchunkWriteOrder.morton:
+            case "morton":
                 subchunk_order_vectorized = _morton_order(chunks_per_shard)
             case _:
                 subchunk_order_vectorized = np.fromiter(

tests/test_codecs/test_sharding.py

@@ -1,6 +1,6 @@
 import pickle
 import re
-from typing import Any
+from typing import Any, get_args
 
 import numpy as np
 import numpy.typing as npt
@@ -562,7 +562,7 @@ def test_sharding_mixed_integer_list_indexing(store: Store) -> None:
 
 @pytest.mark.parametrize(
     "subchunk_write_order",
-    list(SubchunkWriteOrder),
+    get_args(SubchunkWriteOrder),
 )
 async def test_encoded_subchunk_write_order(subchunk_write_order: SubchunkWriteOrder) -> None:
     """Subchunks must be physically laid out in the shard in the order specified by
@@ -611,14 +611,12 @@ async def test_encoded_subchunk_write_order(subchunk_write_order: SubchunkWriteO
     # The physical write order is recovered by sorting coordinates by start offset.
     actual_order = [coord for _, coord in sorted(offset_to_coord.items())]
     expected_order = list(codec._subchunk_order_iter(chunks_per_shard))
-    assert (actual_order == expected_order) == (
-        subchunk_write_order != SubchunkWriteOrder.unordered
-    )
+    assert (actual_order == expected_order) == (subchunk_write_order != "unordered")
 
 
 @pytest.mark.parametrize(
     "subchunk_write_order",
-    list(SubchunkWriteOrder),
+    get_args(SubchunkWriteOrder),
 )
 @pytest.mark.parametrize("do_partial", [True, False], ids=["partial", "complete"])
 def test_subchunk_write_order_roundtrip(