diff --git a/docs/release.rst b/docs/release.rst index 842c36e290..c18e0b8c20 100644 --- a/docs/release.rst +++ b/docs/release.rst @@ -18,6 +18,12 @@ Release notes Unreleased ---------- +Enhancements +~~~~~~~~~~~~ + +* Added type hints to ``zarr.creation.create()``. + By :user:`David Stansby ` :issue:`1536`. + Docs ~~~~ diff --git a/pyproject.toml b/pyproject.toml index 36a0d896ea..4b7fef6003 100644 --- a/pyproject.toml +++ b/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] diff --git a/zarr/_storage/store.py b/zarr/_storage/store.py index 667ca38147..09f0b68602 100644 --- a/zarr/_storage/store.py +++ b/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"] diff --git a/zarr/creation.py b/zarr/creation.py index 6227f90b7b..d4f570895a 100644 --- a/zarr/creation.py +++ b/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. diff --git a/zarr/storage.py b/zarr/storage.py index 5ba8071395..1c3b39862a 100644 --- a/zarr/storage.py +++ b/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] diff --git a/zarr/sync.py b/zarr/sync.py index 49684a51ee..2e843f6557 100644 --- a/zarr/sync.py +++ b/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 `_ package. diff --git a/zarr/tests/test_core.py b/zarr/tests/test_core.py index f3ca73dea8..a3fde4050d 100644 --- a/zarr/tests/test_core.py +++ b/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 diff --git a/zarr/types.py b/zarr/types.py new file mode 100644 index 0000000000..1de270f25c --- /dev/null +++ b/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): + ... diff --git a/zarr/util.py b/zarr/util.py index df1cd9d409..f97094b93a 100644 --- a/zarr/util.py +++ b/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)