From aa8d1e2537cc5fdf035534b3f701b645897f98d1 Mon Sep 17 00:00:00 2001 From: David Stansby Date: Sun, 1 Oct 2023 08:53:49 +0100 Subject: [PATCH 1/8] Add type hints to zarr.create --- zarr/_storage/store.py | 3 ++- zarr/creation.py | 46 +++++++++++++++++++++++------------------ zarr/sync.py | 10 +++++++-- zarr/tests/test_core.py | 8 ++++--- zarr/types.py | 5 +++++ zarr/util.py | 5 +++-- 6 files changed, 49 insertions(+), 28 deletions(-) create mode 100644 zarr/types.py diff --git a/zarr/_storage/store.py b/zarr/_storage/store.py index 8daedae48f..5126aad578 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 726d0b5932..4b74ab250e 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 from zarr.util import normalize_dimension_separator 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. diff --git a/zarr/sync.py b/zarr/sync.py index 49684a51ee..785e0093bb 100644 --- a/zarr/sync.py +++ b/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 `_ 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..42ddd51b44 --- /dev/null +++ b/zarr/types.py @@ -0,0 +1,5 @@ +from typing import Literal + +ZARR_VERSION = Literal[2, 3] +DIMENSION_SEPARATOR = Literal[".", "/"] +MEMORY_ORDER = Literal["C", "F"] diff --git a/zarr/util.py b/zarr/util.py index ea0dd9fcec..3909c29194 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") @@ -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) From ddeabc56868855a693b1c6552812db10e575f635 Mon Sep 17 00:00:00 2001 From: David Stansby Date: Wed, 4 Oct 2023 08:33:18 +0100 Subject: [PATCH 2/8] Use protocol for MetaArray --- zarr/creation.py | 4 ++-- zarr/types.py | 7 ++++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/zarr/creation.py b/zarr/creation.py index 4b74ab250e..38c2022475 100644 --- a/zarr/creation.py +++ b/zarr/creation.py @@ -24,7 +24,7 @@ ) from zarr._storage.store import StorageTransformer from zarr.sync import Synchronizer -from zarr.types import ZARR_VERSION, DIMENSION_SEPARATOR, MEMORY_ORDER +from zarr.types import ZARR_VERSION, DIMENSION_SEPARATOR, MEMORY_ORDER, MetaArray from zarr.util import normalize_dimension_separator @@ -49,7 +49,7 @@ def create( write_empty_chunks: bool = True, *, zarr_version: Optional[ZARR_VERSION] = None, - meta_array: Optional[npt.ArrayLike] = None, + meta_array: Optional[MetaArray] = None, storage_transformers: Sequence[StorageTransformer] = (), **kwargs, ): diff --git a/zarr/types.py b/zarr/types.py index 42ddd51b44..07cd5401d9 100644 --- a/zarr/types.py +++ b/zarr/types.py @@ -1,5 +1,10 @@ -from typing import Literal +from typing import Literal, Protocol ZARR_VERSION = Literal[2, 3] DIMENSION_SEPARATOR = Literal[".", "/"] MEMORY_ORDER = Literal["C", "F"] + + +class MetaArray(Protocol): + def __array_function__(self, func, types, args, kwargs): + ... From 5c17cb03ec7ee36b719e094a852cc87e14adbe2f Mon Sep 17 00:00:00 2001 From: David Stansby Date: Wed, 4 Oct 2023 08:35:43 +0100 Subject: [PATCH 3/8] Use protocol for Synchronizer --- zarr/sync.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/zarr/sync.py b/zarr/sync.py index 785e0093bb..2e843f6557 100644 --- a/zarr/sync.py +++ b/zarr/sync.py @@ -1,14 +1,16 @@ import os from collections import defaultdict from threading import Lock +from typing import Protocol import fasteners -class Synchronizer: +class Synchronizer(Protocol): """Base class for synchronizers.""" - pass + def __getitem__(self, item): + ... class ThreadSynchronizer(Synchronizer): From a60fceb1fc5a4434372a398042da634d4386dd0e Mon Sep 17 00:00:00 2001 From: David Stansby Date: Wed, 4 Oct 2023 08:39:12 +0100 Subject: [PATCH 4/8] Fix Path typing --- zarr/creation.py | 6 +++--- zarr/storage.py | 2 +- zarr/types.py | 5 ++++- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/zarr/creation.py b/zarr/creation.py index 38c2022475..3ea014a3d0 100644 --- a/zarr/creation.py +++ b/zarr/creation.py @@ -24,21 +24,21 @@ ) from zarr._storage.store import StorageTransformer from zarr.sync import Synchronizer -from zarr.types import ZARR_VERSION, DIMENSION_SEPARATOR, MEMORY_ORDER, MetaArray +from zarr.types import ZARR_VERSION, DIMENSION_SEPARATOR, MEMORY_ORDER, MetaArray, PathLike from zarr.util import normalize_dimension_separator def create( shape: Union[int, Tuple[int, ...]], chunks: bool = True, - dtype: Optional[Union[str, npt.DTypeLike]] = None, + dtype: Optional[npt.DTypeLike] = None, compressor="default", fill_value: Optional[int] = 0, order: MEMORY_ORDER = "C", store: Optional[Union[str, MutableMapping]] = None, synchronizer: Optional[Synchronizer] = None, overwrite: bool = False, - path: Optional[str] = None, + path: Optional[PathLike] = None, chunk_store: Optional[MutableMapping] = None, filters: Optional[Sequence[Codec]] = None, cache_metadata: bool = True, diff --git a/zarr/storage.py b/zarr/storage.py index b36f804ebd..9e87eb46dc 100644 --- a/zarr/storage.py +++ b/zarr/storage.py @@ -39,6 +39,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, @@ -104,7 +105,6 @@ default_compressor = Zlib() -Path = Union[str, bytes, None] # allow MutableMapping for backwards compatibility StoreLike = Union[BaseStore, MutableMapping] diff --git a/zarr/types.py b/zarr/types.py index 07cd5401d9..1de270f25c 100644 --- a/zarr/types.py +++ b/zarr/types.py @@ -1,10 +1,13 @@ -from typing import Literal, Protocol +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): ... From 451943a7438cb5c5af35090272c5aa94c01215d2 Mon Sep 17 00:00:00 2001 From: David Stansby Date: Tue, 31 Oct 2023 10:47:24 +0000 Subject: [PATCH 5/8] Add release note --- docs/release.rst | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/release.rst b/docs/release.rst index 2f9b93a361..16e4682eec 100644 --- a/docs/release.rst +++ b/docs/release.rst @@ -18,6 +18,13 @@ Release notes Unreleased ---------- +Enhancements +~~~~~~~~~~~~ + +* Added type hints to ``zarr.creation.create()``. + By :user:`David Stansby ` :issue:`1549`. + + Maintenance ~~~~~~~~~~~ From 424867eabef4006a9029b9f5a14fa86bc3c366ff Mon Sep 17 00:00:00 2001 From: David Stansby Date: Tue, 31 Oct 2023 10:48:44 +0000 Subject: [PATCH 6/8] Fix dim separator typing --- zarr/util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zarr/util.py b/zarr/util.py index 3909c29194..87f6ac2514 100644 --- a/zarr/util.py +++ b/zarr/util.py @@ -289,7 +289,7 @@ def normalize_order(order: str) -> str: def normalize_dimension_separator(sep: Optional[str]) -> Optional[DIMENSION_SEPARATOR]: if sep in (".", "/", None): - return cast(DIMENSION_SEPARATOR, sep) + return cast(Optional[DIMENSION_SEPARATOR], sep) else: raise ValueError("dimension_separator must be either '.' or '/', found: %r" % sep) From 0380cba36f2ac19c31eaffe25de82c4bc8d520fc Mon Sep 17 00:00:00 2001 From: David Stansby Date: Tue, 14 Nov 2023 14:40:11 +0000 Subject: [PATCH 7/8] Ignore ... in coverage reporting --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 6869cbf834..12fd2643b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,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] From 4d22d4f196acd75111183a23ffb5704c7ef6ff01 Mon Sep 17 00:00:00 2001 From: David Stansby Date: Fri, 8 Dec 2023 15:35:07 +0000 Subject: [PATCH 8/8] Fix chunk typing --- zarr/creation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zarr/creation.py b/zarr/creation.py index 3ea014a3d0..a8c0e5229d 100644 --- a/zarr/creation.py +++ b/zarr/creation.py @@ -30,7 +30,7 @@ def create( shape: Union[int, Tuple[int, ...]], - chunks: bool = True, + chunks: Union[int, Tuple[int, ...], bool] = True, dtype: Optional[npt.DTypeLike] = None, compressor="default", fill_value: Optional[int] = 0,