Add type hints to zarr.create

Author: dstansby

zarr/_storage/store.py

@@ -9,6 +9,7 @@
 from zarr.meta import Metadata2, Metadata3
 from zarr.util import normalize_storage_path
 from zarr.context import Context
+from zarr.types import ZARR_VERSION
 
 # v2 store keys
 array_meta_key = ".zarray"
@@ -19,7 +20,7 @@
 meta_root = "meta/root/"
 data_root = "data/root/"
 
-DEFAULT_ZARR_VERSION = 2
+DEFAULT_ZARR_VERSION: ZARR_VERSION = 2
 
 v3_api_available = os.environ.get("ZARR_V3_EXPERIMENTAL_API", "0").lower() not in ["0", "false"]
 

zarr/creation.py

@@ -1,7 +1,9 @@
-from typing import Optional
+from collections.abc import MutableMapping
+from typing import Optional, Tuple, Union, TYPE_CHECKING, Sequence
 from warnings import warn
 
 import numpy as np
+import numpy.typing as npt
 from numcodecs.registry import codec_registry
 
 from zarr._storage.store import DEFAULT_ZARR_VERSION
@@ -19,32 +21,38 @@
     normalize_storage_path,
     normalize_store_arg,
 )
+from zarr.types import ZARR_VERSION, DIMENSION_SEPARATOR, MEMORY_ORDER
 from zarr.util import normalize_dimension_separator
 
+if TYPE_CHECKING:
+    from zarr.sync import Synchronizer
+    from numcodecs.abc import Codec
+    from zarr._storage.store import StorageTransformer
+
 
 def create(
-    shape,
-    chunks=True,
-    dtype=None,
+    shape: Union[int, Tuple[int, ...]],
+    chunks: bool = True,
+    dtype: Optional[Union[str, npt.DTypeLike]] = None,
     compressor="default",
     fill_value: Optional[int] = 0,
-    order="C",
-    store=None,
-    synchronizer=None,
-    overwrite=False,
-    path=None,
-    chunk_store=None,
-    filters=None,
-    cache_metadata=True,
-    cache_attrs=True,
-    read_only=False,
-    object_codec=None,
-    dimension_separator=None,
-    write_empty_chunks=True,
+    order: MEMORY_ORDER = "C",
+    store: Optional[Union[str, MutableMapping]] = None,
+    synchronizer: Optional[Synchronizer] = None,
+    overwrite: bool = False,
+    path: Optional[str] = None,
+    chunk_store: Optional[MutableMapping] = None,
+    filters: Optional[Sequence[Codec]] = None,
+    cache_metadata: bool = True,
+    cache_attrs: bool = True,
+    read_only: bool = False,
+    object_codec: Optional[Codec] = None,
+    dimension_separator: Optional[DIMENSION_SEPARATOR] = None,
+    write_empty_chunks: bool = True,
     *,
-    zarr_version=None,
-    meta_array=None,
-    storage_transformers=(),
+    zarr_version: Optional[ZARR_VERSION] = None,
+    meta_array: Optional[npt.ArrayLike] = None,
+    storage_transformers: Sequence[StorageTransformer] = (),
     **kwargs,
 ):
     """Create an array.

zarr/sync.py

@@ -5,7 +5,13 @@
 import fasteners
 
 
-class ThreadSynchronizer:
+class Synchronizer:
+    """Base class for synchronizers."""
+
+    pass
+
+
+class ThreadSynchronizer(Synchronizer):
     """Provides synchronization using thread locks."""
 
     def __init__(self):
@@ -24,7 +30,7 @@ def __setstate__(self, *args):
         self.__init__()
 
 
-class ProcessSynchronizer:
+class ProcessSynchronizer(Synchronizer):
     """Provides synchronization using file locks via the
     `fasteners <https://fasteners.readthedocs.io/en/latest/api/inter_process/>`_
     package.

zarr/tests/test_core.py

@@ -3,7 +3,7 @@
 import sys
 import pickle
 import shutil
-from typing import Any, Literal, Optional, Tuple, Union
+from typing import Any, Literal, Optional, Tuple, Union, Sequence
 import unittest
 from itertools import zip_longest
 from tempfile import mkdtemp
@@ -26,6 +26,7 @@
     VLenUTF8,
     Zlib,
 )
+from numcodecs.abc import Codec
 from numcodecs.compat import ensure_bytes, ensure_ndarray
 from numcodecs.tests.common import greetings
 from numpy.testing import assert_array_almost_equal, assert_array_equal
@@ -73,6 +74,7 @@
 from zarr.tests.test_storage_v3 import DummyStorageTransfomer
 from zarr.util import buffer_size
 from zarr.tests.util import abs_container, skip_test_env_var, have_fsspec, mktemp
+from zarr.types import DIMENSION_SEPARATOR
 
 # noinspection PyMethodMayBeStatic
 
@@ -82,8 +84,8 @@ class TestArray:
     root = ""
     path = ""
     compressor = Zlib(level=1)
-    filters = None
-    dimension_separator: Literal["/", ".", None] = None
+    filters: Optional[Sequence[Codec]] = None
+    dimension_separator: Optional[DIMENSION_SEPARATOR] = None
     cache_metadata = True
     cache_attrs = True
     partial_decompress: bool = False

zarr/types.py

@@ -0,0 +1,5 @@
+from typing import Literal
+
+ZARR_VERSION = Literal[2, 3]
+DIMENSION_SEPARATOR = Literal[".", "/"]
+MEMORY_ORDER = Literal["C", "F"]

zarr/util.py

@@ -31,6 +31,7 @@
 from numcodecs.ndarray_like import NDArrayLike
 from numcodecs.registry import codec_registry
 from numcodecs.blosc import cbuffer_sizes, cbuffer_metainfo
+from zarr.types import DIMENSION_SEPARATOR
 
 KeyType = TypeVar("KeyType")
 ValueType = TypeVar("ValueType")
@@ -286,9 +287,9 @@ def normalize_order(order: str) -> str:
     return order
 
 
-def normalize_dimension_separator(sep: Optional[str]) -> Optional[str]:
+def normalize_dimension_separator(sep: Optional[str]) -> Optional[DIMENSION_SEPARATOR]:
     if sep in (".", "/", None):
-        return sep
+        return cast(DIMENSION_SEPARATOR, sep)
     else:
         raise ValueError("dimension_separator must be either '.' or '/', found: %r" % sep)