fix: attempt to resolve typing issues except covariance

Author: ilan-gold

src/zarr/_compat.py

@@ -2,16 +2,14 @@
 from collections.abc import Callable
 from functools import wraps
 from inspect import Parameter, signature
-from typing import Any, TypeVar
+from typing import Any
 
 from zarr.errors import ZarrFutureWarning
 
-T = TypeVar("T")
-
 # Based off https://github.com/scikit-learn/scikit-learn/blob/e87b32a81c70abed8f2e97483758eb64df8255e9/sklearn/utils/validation.py#L63
 
 
-def _deprecate_positional_args(
+def _deprecate_positional_args[T](
     func: Callable[..., T] | None = None, *, version: str = "3.1.0"
 ) -> Callable[..., T]:
     """Decorator for methods that issues warnings for positional arguments.

src/zarr/abc/codec.py

@@ -2,7 +2,7 @@
 
 from abc import abstractmethod
 from collections.abc import Mapping
-from typing import TYPE_CHECKING, Generic, TypeGuard, TypeVar
+from typing import TYPE_CHECKING, TypeGuard
 
 from typing_extensions import ReadOnly, TypedDict
 
@@ -34,13 +34,11 @@
     "CodecPipeline",
 ]
 
-CodecInput = TypeVar("CodecInput", bound=NDBuffer | Buffer)
-CodecOutput = TypeVar("CodecOutput", bound=NDBuffer | Buffer)
+type CodecInput = NDBuffer | Buffer
+type CodecOutput = NDBuffer | Buffer
 
-TName = TypeVar("TName", bound=str, covariant=True)
 
-
-class CodecJSON_V2(TypedDict, Generic[TName]):
+class CodecJSON_V2[TName: str](TypedDict):
     """The JSON representation of a codec for Zarr V2"""
 
     id: ReadOnly[TName]
@@ -59,7 +57,7 @@ def _check_codecjson_v2(data: object) -> TypeGuard[CodecJSON_V2[str]]:
 """The widest type of JSON-like input that could specify a codec."""
 
 
-class BaseCodec(Metadata, Generic[CodecInput, CodecOutput]):
+class BaseCodec[CI: CodecInput, CO: CodecOutput](Metadata):
     """Generic base class for codecs.
 
     Codecs can be registered via zarr.codecs.registry.
@@ -137,13 +135,13 @@ def validate(
             The array chunk grid
         """
 
-    async def _decode_single(self, chunk_data: CodecOutput, chunk_spec: ArraySpec) -> CodecInput:
+    async def _decode_single(self, chunk_data: CO, chunk_spec: ArraySpec) -> CI:
         raise NotImplementedError  # pragma: no cover
 
     async def decode(
         self,
-        chunks_and_specs: Iterable[tuple[CodecOutput | None, ArraySpec]],
-    ) -> Iterable[CodecInput | None]:
+        chunks_and_specs: Iterable[tuple[CO | None, ArraySpec]],
+    ) -> Iterable[CI | None]:
         """Decodes a batch of chunks.
         Chunks can be None in which case they are ignored by the codec.
 
@@ -154,25 +152,23 @@ async def decode(
 
         Returns
         -------
-        Iterable[CodecInput | None]
+        Iterable[CI | None]
         """
         return await _batching_helper(self._decode_single, chunks_and_specs)
 
-    async def _encode_single(
-        self, chunk_data: CodecInput, chunk_spec: ArraySpec
-    ) -> CodecOutput | None:
+    async def _encode_single(self, chunk_data: CI, chunk_spec: ArraySpec) -> CO | None:
         raise NotImplementedError  # pragma: no cover
 
     async def encode(
         self,
-        chunks_and_specs: Iterable[tuple[CodecInput | None, ArraySpec]],
-    ) -> Iterable[CodecOutput | None]:
+        chunks_and_specs: Iterable[tuple[CI | None, ArraySpec]],
+    ) -> Iterable[CO | None]:
         """Encodes a batch of chunks.
         Chunks can be None in which case they are ignored by the codec.
 
         Parameters
         ----------
-        chunks_and_specs : Iterable[tuple[CodecInput | None, ArraySpec]]
+        chunks_and_specs : Iterable[tuple[CI | None, ArraySpec]]
             Ordered set of to-be-encoded chunks with their accompanying chunk spec.
 
         Returns
@@ -460,21 +456,21 @@ async def write(
         ...
 
 
-async def _batching_helper(
-    func: Callable[[CodecInput, ArraySpec], Awaitable[CodecOutput | None]],
-    batch_info: Iterable[tuple[CodecInput | None, ArraySpec]],
-) -> list[CodecOutput | None]:
+async def _batching_helper[CI: CodecInput, CO: CodecOutput](
+    func: Callable[[CI, ArraySpec], Awaitable[CO | None]],
+    batch_info: Iterable[tuple[CI | None, ArraySpec]],
+) -> list[CO | None]:
     return await concurrent_map(
         list(batch_info),
         _noop_for_none(func),
         config.get("async.concurrency"),
     )
 
 
-def _noop_for_none(
-    func: Callable[[CodecInput, ArraySpec], Awaitable[CodecOutput | None]],
-) -> Callable[[CodecInput | None, ArraySpec], Awaitable[CodecOutput | None]]:
-    async def wrap(chunk: CodecInput | None, chunk_spec: ArraySpec) -> CodecOutput | None:
+def _noop_for_none[CI: CodecInput, CO: CodecOutput](
+    func: Callable[[CI, ArraySpec], Awaitable[CO | None]],
+) -> Callable[[CI | None, ArraySpec], Awaitable[CO | None]]:
+    async def wrap(chunk: CI | None, chunk_spec: ArraySpec) -> CO | None:
         if chunk is None:
             return None
         return await func(chunk, chunk_spec)

src/zarr/abc/numcodec.py

@@ -1,6 +1,4 @@
-from typing import Any, Self, TypeGuard
-
-from typing_extensions import Protocol
+from typing import Any, Protocol, Self, TypeGuard
 
 
 class Numcodec(Protocol):

src/zarr/abc/store.py

@@ -9,7 +9,7 @@
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, AsyncIterator, Iterable
     from types import TracebackType
-    from typing import Any, Self, TypeAlias
+    from typing import Any, Self
 
     from zarr.core.buffer import Buffer, BufferPrototype
 
@@ -42,7 +42,7 @@ class SuffixByteRequest:
     """The number of bytes from the suffix to request."""
 
 
-ByteRequest: TypeAlias = RangeByteRequest | OffsetByteRequest | SuffixByteRequest
+type ByteRequest = RangeByteRequest | OffsetByteRequest | SuffixByteRequest
 
 
 class Store(ABC):

src/zarr/core/array.py

@@ -10,9 +10,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
-    Generic,
     Literal,
-    TypeAlias,
     TypedDict,
     cast,
     overload,
@@ -107,7 +105,6 @@
     ArrayV2Metadata,
     ArrayV2MetadataDict,
     ArrayV3Metadata,
-    T_ArrayMetadata,
 )
 from zarr.core.metadata.io import save_metadata
 from zarr.core.metadata.v2 import (
@@ -278,7 +275,7 @@ async def get_array_metadata(
 
 
 @dataclass(frozen=True)
-class AsyncArray(Generic[T_ArrayMetadata]):
+class AsyncArray[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]:
     """
     An asynchronous array class representing a chunked array stored in a Zarr store.
 
@@ -4246,7 +4243,7 @@ async def _shards_initialized(
     )
 
 
-FiltersLike: TypeAlias = (
+type FiltersLike = (
     Iterable[dict[str, JSON] | ArrayArrayCodec | Numcodec]
     | ArrayArrayCodec
     | Iterable[Numcodec]
@@ -4255,25 +4252,25 @@ async def _shards_initialized(
     | None
 )
 # Union of acceptable types for users to pass in for both v2 and v3 compressors
-CompressorLike: TypeAlias = dict[str, JSON] | BytesBytesCodec | Numcodec | Literal["auto"] | None
+type CompressorLike = dict[str, JSON] | BytesBytesCodec | Numcodec | Literal["auto"] | None
 
-CompressorsLike: TypeAlias = (
+type CompressorsLike = (
     Iterable[dict[str, JSON] | BytesBytesCodec | Numcodec]
     | Mapping[str, JSON]
     | BytesBytesCodec
     | Numcodec
     | Literal["auto"]
     | None
 )
-SerializerLike: TypeAlias = dict[str, JSON] | ArrayBytesCodec | Literal["auto"]
+type SerializerLike = dict[str, JSON] | ArrayBytesCodec | Literal["auto"]
 
 
 class ShardsConfigParam(TypedDict):
     shape: tuple[int, ...]
     index_location: ShardingCodecIndexLocation | None
 
 
-ShardsLike: TypeAlias = tuple[int, ...] | ShardsConfigParam | Literal["auto"]
+type ShardsLike = tuple[int, ...] | ShardsConfigParam | Literal["auto"]
 
 
 async def from_array(

src/zarr/core/buffer/core.py

@@ -267,7 +267,7 @@ def as_buffer_like(self) -> BytesLike:
         -------
             An object that implements the Python buffer protocol
         """
-        return memoryview(self.as_numpy_array())  # type: ignore[arg-type]
+        return memoryview(self.as_numpy_array())
 
     def to_bytes(self) -> bytes:
         """Returns the buffer as `bytes` (host memory).

src/zarr/core/chunk_key_encodings.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypedDict, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, cast
 
 if TYPE_CHECKING:
     from typing import NotRequired, Self
@@ -62,7 +62,7 @@ def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str:
         """
 
 
-ChunkKeyEncodingLike: TypeAlias = (
+type ChunkKeyEncodingLike = (
     dict[str, JSON] | ChunkKeyEncodingParams | ChunkKeyEncoding | NamedConfig[str, Any]
 )
 

src/zarr/core/codec_pipeline.py

@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass
 from itertools import islice, pairwise
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import TYPE_CHECKING, Any
 from warnings import warn
 
 from zarr.abc.codec import (
@@ -30,11 +30,8 @@
     from zarr.core.chunk_grids import ChunkGrid
     from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType
 
-T = TypeVar("T")
-U = TypeVar("U")
 
-
-def _unzip2(iterable: Iterable[tuple[T, U]]) -> tuple[list[T], list[U]]:
+def _unzip2[T, U](iterable: Iterable[tuple[T, U]]) -> tuple[list[T], list[U]]:
     out0: list[T] = []
     out1: list[U] = []
     for item0, item1 in iterable:
@@ -43,7 +40,7 @@ def _unzip2(iterable: Iterable[tuple[T, U]]) -> tuple[list[T], list[U]]:
     return (out0, out1)
 
 
-def batched(iterable: Iterable[T], n: int) -> Iterable[tuple[T, ...]]:
+def batched[T](iterable: Iterable[T], n: int) -> Iterable[tuple[T, ...]]:
     if n < 1:
         raise ValueError("n must be at least one")
     it = iter(iterable)

src/zarr/core/common.py

@@ -12,11 +12,9 @@
     TYPE_CHECKING,
     Any,
     Final,
-    Generic,
     Literal,
     NotRequired,
     TypedDict,
-    TypeVar,
     cast,
     overload,
 )
@@ -48,11 +46,8 @@
 ANY_ACCESS_MODE: Final = "r", "r+", "a", "w", "w-"
 DimensionNames = Iterable[str | None] | None
 
-TName = TypeVar("TName", bound=str)
-TConfig = TypeVar("TConfig", bound=Mapping[str, object])
 
-
-class NamedConfig(TypedDict, Generic[TName, TConfig]):
+class NamedConfig[TName: str, TConfig: Mapping[str, object]](TypedDict):
     """
     A typed dictionary representing an object with a name and configuration, where the configuration
     is an optional mapping of string keys to values, e.g. another typed dictionary or a JSON object.
@@ -68,7 +63,7 @@ class NamedConfig(TypedDict, Generic[TName, TConfig]):
     """The configuration of the object. Not required."""
 
 
-class NamedRequiredConfig(TypedDict, Generic[TName, TConfig]):
+class NamedRequiredConfig[TName: str, TConfig: Mapping[str, object]](TypedDict):
     """
     A typed dictionary representing an object with a name and configuration, where the configuration
     is a mapping of string keys to values, e.g. another typed dictionary or a JSON object.
@@ -94,11 +89,7 @@ def ceildiv(a: float, b: float) -> int:
     return math.ceil(a / b)
 
 
-T = TypeVar("T", bound=tuple[Any, ...])
-V = TypeVar("V")
-
-
-async def concurrent_map(
+async def concurrent_map[T: tuple[Any, ...], V](
     items: Iterable[T],
     func: Callable[..., Awaitable[V]],
     limit: int | None = None,
@@ -116,15 +107,12 @@ async def run(item: tuple[Any]) -> V:
         return await asyncio.gather(*[asyncio.ensure_future(run(item)) for item in items])
 
 
-E = TypeVar("E", bound=Enum)
-
-
-def enum_names(enum: type[E]) -> Iterator[str]:
+def enum_names[E: Enum](enum: type[E]) -> Iterator[str]:
     for item in enum:
         yield item.name
 
 
-def parse_enum(data: object, cls: type[E]) -> E:
+def parse_enum[E: Enum](data: object, cls: type[E]) -> E:
     if isinstance(data, cls):
         return data
     if not isinstance(data, str):

src/zarr/core/dtype/__init__.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Final, TypeAlias
+from typing import TYPE_CHECKING, Final
 
 from zarr.core.dtype.common import (
     DataTypeValidationError,
@@ -150,14 +150,15 @@
 VLEN_UTF8_ALIAS: Final = ("str", str, "string")
 
 # This type models inputs that can be coerced to a ZDType
-ZDTypeLike: TypeAlias = npt.DTypeLike | ZDType[TBaseDType, TBaseScalar] | Mapping[str, JSON] | str
+type ZDTypeLike = npt.DTypeLike | ZDType[TBaseDType, TBaseScalar] | Mapping[str, JSON] | str
 
 for dtype in ANY_DTYPE:
     # mypy does not know that all the elements of ANY_DTYPE are subclasses of ZDType
     data_type_registry.register(dtype._zarr_v3_name, dtype)  # type: ignore[arg-type]
 
 
 # TODO: find a better name for this function
+# TODO: Can we type this better? Otherwise functions that use it can't be specific enough about TBaseDType
 def get_data_type_from_native_dtype(dtype: npt.DTypeLike) -> ZDType[TBaseDType, TBaseScalar]:
     """
     Get a data type wrapper (an instance of ``ZDType``) from a native data type, e.g. a numpy dtype.