33import asyncio
44from collections .abc import Mapping
55from dataclasses import dataclass , replace
6+ from enum import Enum
67from functools import cached_property
78from typing import (
89 TYPE_CHECKING ,
3435 from zarr .core .array_spec import ArraySpec
3536 from zarr .core .buffer import Buffer
3637
37- BloscShuffle = Literal ["noshuffle" , "shuffle" , "bitshuffle" ]
38+
39+ class BloscShuffle (Enum ):
40+ """
41+ Enum for shuffle filter used by blosc.
42+ """
43+
44+ noshuffle = "noshuffle"
45+ shuffle = "shuffle"
46+ bitshuffle = "bitshuffle"
47+
48+ @classmethod
49+ def from_int (cls , num : int ) -> BloscShuffle :
50+ blosc_shuffle_int_to_str = {
51+ 0 : "noshuffle" ,
52+ 1 : "shuffle" ,
53+ 2 : "bitshuffle" ,
54+ }
55+ if num not in blosc_shuffle_int_to_str :
56+ raise ValueError (f"Value must be between 0 and 2. Got { num } )
57+ return BloscShuffle [blosc_shuffle_int_to_str [num ]]
58+
59+
60+ class BloscCname (Enum ):
61+ """
62+ Enum for compression library used by blosc.
63+ """
64+
65+ lz4 = "lz4"
66+ lz4hc = "lz4hc"
67+ blosclz = "blosclz"
68+ zstd = "zstd"
69+ snappy = "snappy"
70+ zlib = "zlib"
71+
72+
73+ # TODO: Rename this when we retire the enums
74+ BloscShuffle_lit = Literal ["noshuffle" , "shuffle" , "bitshuffle" ]
75+ """The names of the shuffle options used by the blosc codec."""
3876BLOSC_SHUFFLE : Final = ("noshuffle" , "shuffle" , "bitshuffle" )
3977
40- BloscCname = Literal ["lz4" , "lz4hc" , "blosclz" , "zstd" , "snappy" , "zlib" ]
78+ # TODO: rename this when we retire the enums
79+ BloscCname_lit = Literal ["lz4" , "lz4hc" , "blosclz" , "zstd" , "snappy" , "zlib" ]
80+ """The names of the compression libraries used by the blosc codec"""
4181BLOSC_CNAME : Final = ("lz4" , "lz4hc" , "blosclz" , "zstd" , "snappy" , "zlib" )
4282
4383
4484class BloscConfigV2 (TypedDict ):
45- cname : BloscCname
85+ cname : BloscCname_lit
4686 clevel : int
4787 shuffle : int
4888 blocksize : int
4989 typesize : NotRequired [int ]
5090
5191
5292class BloscConfigV3 (TypedDict ):
53- cname : BloscCname
93+ cname : BloscCname_lit
5494 clevel : int
55- shuffle : BloscShuffle
95+ shuffle : BloscShuffle_lit
5696 blocksize : int
5797 typesize : int
5898
@@ -91,9 +131,14 @@ def check_json_v3(data: object) -> TypeGuard[BloscJSON_V3]:
91131
92132
93133def parse_cname (value : object ) -> BloscCname :
94- if value not in BLOSC_CNAME :
95- raise ValueError (f"Value must be one of { BLOSC_CNAME } { value } )
96- return value # type: ignore[return-value]
134+ if isinstance (value , BloscCname ):
135+ return value
136+ if isinstance (value , str ):
137+ if value not in BLOSC_CNAME :
138+ raise ValueError (f"Value must be one of { BLOSC_CNAME } { value } )
139+ return BloscCname [value ]
140+ msg = f"Value must be an instance of `BloscCname` or a string in { BLOSC_CNAME } { value }
141+ raise TypeError (msg )
97142
98143
99144# See https://zarr.readthedocs.io/en/stable/user-guide/performance.html#configuring-blosc
@@ -125,8 +170,10 @@ def parse_blocksize(data: JSON) -> int:
125170
126171
127172def parse_shuffle (data : object ) -> BloscShuffle :
173+ if isinstance (data , BloscShuffle ):
174+ return data
128175 if data in BLOSC_SHUFFLE :
129- return data # type: ignore[return-value ]
176+ return BloscShuffle [ data ] # type: ignore[misc ]
130177 raise TypeError (f"Value must be one of { BLOSC_SHUFFLE } { data } )
131178
132179
@@ -144,9 +191,9 @@ def __init__(
144191 self ,
145192 * ,
146193 typesize : int | None = None ,
147- cname : BloscCname = "zstd" ,
194+ cname : BloscCname_lit | BloscCname = "zstd" ,
148195 clevel : int = 5 ,
149- shuffle : BloscShuffle | None = None ,
196+ shuffle : BloscShuffle_lit | BloscShuffle | None = None ,
150197 blocksize : int = 0 ,
151198 ) -> None :
152199 typesize_parsed = parse_typesize (typesize ) if typesize is not None else None
@@ -174,8 +221,8 @@ def to_dict(self) -> dict[str, JSON]:
174221 "name" : "blosc" ,
175222 "configuration" : {
176223 "clevel" : self .clevel ,
177- "cname" : self .cname ,
178- "shuffle" : self .shuffle ,
224+ "cname" : self .cname . value ,
225+ "shuffle" : self .shuffle . value ,
179226 "typesize" : self .typesize ,
180227 "blocksize" : self .blocksize ,
181228 },
@@ -226,8 +273,8 @@ def to_json(self, zarr_format: ZarrFormat) -> BloscJSON_V2 | BloscJSON_V3:
226273 return {
227274 "id" : "blosc" ,
228275 "clevel" : self .clevel ,
229- "cname" : self .cname ,
230- "shuffle" : BLOSC_SHUFFLE .index (self .shuffle ),
276+ "cname" : self .cname . value ,
277+ "shuffle" : BLOSC_SHUFFLE .index (self .shuffle . value ),
231278 "blocksize" : self .blocksize ,
232279 }
233280 elif zarr_format == 3 :
@@ -244,7 +291,10 @@ def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
244291 if new_codec .typesize is None :
245292 new_codec = replace (new_codec , typesize = item_size )
246293 if new_codec .shuffle is None :
247- new_codec = replace (new_codec , shuffle = "bitshuffle" if item_size == 1 else "shuffle" )
294+ new_codec = replace (
295+ new_codec ,
296+ shuffle = BloscShuffle .bitshuffle if item_size == 1 else BloscShuffle .shuffle ,
297+ )
248298
249299 return new_codec
250300
@@ -255,9 +305,9 @@ def _blosc_codec(self) -> Blosc:
255305 if self .typesize is None :
256306 raise ValueError ("`typesize` needs to be set for decoding and encoding." )
257307 config_dict = {
258- "cname" : self .cname ,
308+ "cname" : self .cname . value ,
259309 "clevel" : self .clevel ,
260- "shuffle" : BLOSC_SHUFFLE .index (self .shuffle ),
310+ "shuffle" : BLOSC_SHUFFLE .index (self .shuffle . value ),
261311 "blocksize" : self .blocksize ,
262312 }
263313 # See https://github.com/zarr-developers/numcodecs/pull/713
0 commit comments