Add type hints to zarr.create (#1536)

* Add type hints to zarr.create

* Use protocol for MetaArray

* Use protocol for Synchronizer

* Fix Path typing

* Add release note

* Fix dim separator typing

* Ignore ... in coverage reporting

* Fix chunk typing

---------

Co-authored-by: Davis Bennett <[email protected]>

Author: dstansby

docs/release.rst

@@ -18,6 +18,12 @@ Release notes
 Unreleased
 ----------
 
+Enhancements
+~~~~~~~~~~~~
+
+* Added type hints to ``zarr.creation.create()``.
+  By :user:`David Stansby <dstansby>` :issue:`1536`.
+
 Docs
 ~~~~
 

pyproject.toml

@@ -65,6 +65,7 @@ Homepage = "https://github.com/zarr-developers/zarr-python"
 exclude_lines = [
     "pragma: no cover",
     "pragma: ${PY_MAJOR_VERSION} no cover",
+    '.*\.\.\.' # Ignore "..." lines
 ]
 
 [tool.coverage.run]

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,10 @@
-from typing import Optional
+from collections.abc import MutableMapping
+from typing import Optional, Tuple, Union, Sequence
 from warnings import warn
 
 import numpy as np
+import numpy.typing as npt
+from numcodecs.abc import Codec
 from numcodecs.registry import codec_registry
 
 from zarr._storage.store import DEFAULT_ZARR_VERSION
@@ -19,32 +22,35 @@
     normalize_storage_path,
     normalize_store_arg,
 )
+from zarr._storage.store import StorageTransformer
+from zarr.sync import Synchronizer
+from zarr.types import ZARR_VERSION, DIMENSION_SEPARATOR, MEMORY_ORDER, MetaArray, PathLike
 from zarr.util import normalize_dimension_separator
 
 
 def create(
-    shape,
-    chunks=True,
-    dtype=None,
+    shape: Union[int, Tuple[int, ...]],
+    chunks: Union[int, Tuple[int, ...], bool] = True,
+    dtype: Optional[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[PathLike] = 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[MetaArray] = None,
+    storage_transformers: Sequence[StorageTransformer] = (),
     **kwargs,
 ):
     """Create an array.

zarr/storage.py

@@ -40,6 +40,7 @@
 from numcodecs.compat import ensure_bytes, ensure_text, ensure_contiguous_ndarray_like
 from numcodecs.registry import codec_registry
 from zarr.context import Context
+from zarr.types import PathLike as Path
 
 from zarr.errors import (
     MetadataError,
@@ -105,7 +106,6 @@
     default_compressor = Zlib()
 
 
-Path = Union[str, bytes, None]
 # allow MutableMapping for backwards compatibility
 StoreLike = Union[BaseStore, MutableMapping]
 

zarr/sync.py

@@ -1,11 +1,19 @@
 import os
 from collections import defaultdict
 from threading import Lock
+from typing import Protocol
 
 import fasteners
 
 
-class ThreadSynchronizer:
+class Synchronizer(Protocol):
+    """Base class for synchronizers."""
+
+    def __getitem__(self, item):
+        ...
+
+
+class ThreadSynchronizer(Synchronizer):
     """Provides synchronization using thread locks."""
 
     def __init__(self):
@@ -24,7 +32,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,13 @@
+from typing import Literal, Protocol, Union
+
+ZARR_VERSION = Literal[2, 3]
+DIMENSION_SEPARATOR = Literal[".", "/"]
+MEMORY_ORDER = Literal["C", "F"]
+
+
+PathLike = Union[str, bytes, None]
+
+
+class MetaArray(Protocol):
+    def __array_function__(self, func, types, args, kwargs):
+        ...

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")
@@ -284,9 +285,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(Optional[DIMENSION_SEPARATOR], sep)
     else:
         raise ValueError("dimension_separator must be either '.' or '/', found: %r" % sep)