fix: try covariance handling

Author: ilan-gold

foo.py

@@ -0,0 +1,17 @@
+class Base1: ...
+class SubBase1(Base1): ...
+class SubSubBase1(SubBase1): ...
+
+class Base2: ...
+class SubBase2(Base2): ...
+class SubSubBase2(SubBase2): ...
+
+class F[T: Base1, S: Base2]: ...
+
+class FSub(F[SubBase1, SubBase2]): ...
+
+class FSubSub(F[SubSubBase1, SubSubBase2]): ...
+
+type GType = F
+
+g: GType = FSubSub()

src/zarr/core/dtype/npy/complex.py

@@ -22,8 +22,6 @@
 )
 from zarr.core.dtype.npy.common import (
     ComplexLike,
-    TComplexDType_co,
-    TComplexScalar_co,
     check_json_complex_float_v2,
     check_json_complex_float_v3,
     complex_float_from_json_v2,
@@ -40,7 +38,10 @@
 
 
 @dataclass(frozen=True)
-class BaseComplex(ZDType[TComplexDType_co, TComplexScalar_co], HasEndianness, HasItemSize):
+class BaseComplex[
+    DType: np.dtypes.Complex64DType | np.dtypes.Complex128DType,
+    Scalar: np.complex64 | np.complex128,
+](ZDType[DType, Scalar], HasEndianness, HasItemSize):
     """
     A base class for Zarr data types that wrap NumPy complex float data types.
     """
@@ -74,13 +75,13 @@ def from_native_dtype(cls, dtype: TBaseDType) -> Self:
             f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}"
         )
 
-    def to_native_dtype(self) -> TComplexDType_co:
+    def to_native_dtype(self) -> DType:
         """
         Convert this class to a NumPy complex dtype with the appropriate byte order.
 
         Returns
         -------
-        TComplexDType_co
+        DType
             A NumPy data type object representing the complex data type with the specified byte order.
         """
 
@@ -235,7 +236,7 @@ def _check_scalar(self, data: object) -> TypeGuard[ComplexLike]:
         """
         return isinstance(data, ComplexLike)
 
-    def _cast_scalar_unchecked(self, data: ComplexLike) -> TComplexScalar_co:
+    def _cast_scalar_unchecked(self, data: ComplexLike) -> Scalar:
         """
         Cast the provided scalar data to the native scalar type of this class.
 
@@ -246,7 +247,7 @@ def _cast_scalar_unchecked(self, data: ComplexLike) -> TComplexScalar_co:
 
         Returns
         -------
-        TComplexScalar_co
+        Scalar
             The casted data as a numpy complex scalar.
 
         Notes
@@ -256,7 +257,7 @@ def _cast_scalar_unchecked(self, data: ComplexLike) -> TComplexScalar_co:
         """
         return self.to_native_dtype().type(data)  # type: ignore[return-value]
 
-    def cast_scalar(self, data: object) -> TComplexScalar_co:
+    def cast_scalar(self, data: object) -> Scalar:
         """
         Attempt to cast a given object to a numpy complex scalar.
 
@@ -267,7 +268,7 @@ def cast_scalar(self, data: object) -> TComplexScalar_co:
 
         Returns
         -------
-        TComplexScalar_co
+        Scalar
             The data cast as a numpy complex scalar.
 
         Raises
@@ -283,7 +284,7 @@ def cast_scalar(self, data: object) -> TComplexScalar_co:
         )
         raise TypeError(msg)
 
-    def default_scalar(self) -> TComplexScalar_co:
+    def default_scalar(self) -> Scalar:
         """
         Get the default value, which is 0 cast to this dtype
 
@@ -294,7 +295,7 @@ def default_scalar(self) -> TComplexScalar_co:
         """
         return self._cast_scalar_unchecked(0)
 
-    def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> TComplexScalar_co:
+    def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> Scalar:
         """
         Read a JSON-serializable value as a numpy float.
 

src/zarr/core/dtype/npy/float.py

@@ -15,8 +15,6 @@
 )
 from zarr.core.dtype.npy.common import (
     FloatLike,
-    TFloatDType_co,
-    TFloatScalar_co,
     check_json_float_v2,
     check_json_float_v3,
     check_json_floatish_str,
@@ -34,7 +32,10 @@
 
 
 @dataclass(frozen=True)
-class BaseFloat(ZDType[TFloatDType_co, TFloatScalar_co], HasEndianness, HasItemSize):
+class BaseFloat[
+    DType: np.dtypes.Float16DType | np.dtypes.Float32DType | np.dtypes.Float64DType,
+    Scalar: np.float16 | np.float32 | np.float64,
+](ZDType[DType, Scalar], HasEndianness, HasItemSize):
     """
     A base class for Zarr data types that wrap NumPy float data types.
     """
@@ -63,13 +64,13 @@ def from_native_dtype(cls, dtype: TBaseDType) -> Self:
             f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}"
         )
 
-    def to_native_dtype(self) -> TFloatDType_co:
+    def to_native_dtype(self) -> DType:
         """
         Convert the wrapped data type to a NumPy data type.
 
         Returns
         -------
-        TFloatDType_co
+        DType
             The NumPy data type.
         """
         byte_order = endianness_to_numpy_str(self.endianness)
@@ -203,7 +204,7 @@ def _check_scalar(self, data: object) -> TypeGuard[FloatLike]:
         """
         return isinstance(data, FloatLike)
 
-    def _cast_scalar_unchecked(self, data: FloatLike) -> TFloatScalar_co:
+    def _cast_scalar_unchecked(self, data: FloatLike) -> Scalar:
         """
         Cast a scalar value to a NumPy float scalar.
 
@@ -214,12 +215,12 @@ def _cast_scalar_unchecked(self, data: FloatLike) -> TFloatScalar_co:
 
         Returns
         -------
-        TFloatScalar_co
+        Scalar
             The NumPy float scalar.
         """
         return self.to_native_dtype().type(data)  # type: ignore[return-value]
 
-    def cast_scalar(self, data: object) -> TFloatScalar_co:
+    def cast_scalar(self, data: object) -> Scalar:
         """
         Cast a scalar value to a NumPy float scalar.
 
@@ -230,7 +231,7 @@ def cast_scalar(self, data: object) -> TFloatScalar_co:
 
         Returns
         -------
-        TFloatScalar_co
+        Scalar
             The NumPy float scalar.
         """
         if self._check_scalar(data):
@@ -241,18 +242,18 @@ def cast_scalar(self, data: object) -> TFloatScalar_co:
         )
         raise TypeError(msg)
 
-    def default_scalar(self) -> TFloatScalar_co:
+    def default_scalar(self) -> Scalar:
         """
         Get the default value, which is 0 cast to this zdtype.
 
         Returns
         -------
-        TFloatScalar_co
+        Scalar
             The default value.
         """
         return self._cast_scalar_unchecked(0)
 
-    def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> TFloatScalar_co:
+    def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> Scalar:
         """
         Read a JSON-serializable value as a NumPy float scalar.
 
@@ -265,7 +266,7 @@ def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> TFloatScal
 
         Returns
         -------
-        TFloatScalar_co
+        Scalar
             The NumPy float scalar.
         """
         if zarr_format == 2:

src/zarr/core/dtype/npy/int.py

@@ -54,7 +54,10 @@
 
 
 @dataclass(frozen=True)
-class BaseInt(ZDType[TIntDType_co, TIntScalar_co], HasItemSize):
+class BaseInt[
+    DType: _NumpyIntDType,
+    Scalar: np.int8 | np.int16 | np.int32 | np.int64 | np.uint8 | np.uint16 | np.uint32 | np.uint64,
+](ZDType[DType, Scalar], HasItemSize):
     """
     A base class for integer data types in Zarr.
 
@@ -129,7 +132,7 @@ def _check_scalar(self, data: object) -> TypeGuard[IntLike]:
 
         return isinstance(data, IntLike)
 
-    def _cast_scalar_unchecked(self, data: IntLike) -> TIntScalar_co:
+    def _cast_scalar_unchecked(self, data: IntLike) -> Scalar:
         """
         Casts a given scalar value to the native integer scalar type without type checking.
 
@@ -140,13 +143,13 @@ def _cast_scalar_unchecked(self, data: IntLike) -> TIntScalar_co:
 
         Returns
         -------
-        TIntScalar_co
+        Scalar
             The casted integer scalar of the native dtype.
         """
 
         return self.to_native_dtype().type(data)  # type: ignore[return-value]
 
-    def cast_scalar(self, data: object) -> TIntScalar_co:
+    def cast_scalar(self, data: object) -> Scalar:
         """
         Attempt to cast a given object to a NumPy integer scalar.
 
@@ -157,7 +160,7 @@ def cast_scalar(self, data: object) -> TIntScalar_co:
 
         Returns
         -------
-        TIntScalar_co
+        Scalar
             The data cast as a NumPy integer scalar.
 
         Raises
@@ -174,18 +177,18 @@ def cast_scalar(self, data: object) -> TIntScalar_co:
         )
         raise TypeError(msg)
 
-    def default_scalar(self) -> TIntScalar_co:
+    def default_scalar(self) -> Scalar:
         """
         Get the default value, which is 0 cast to this dtype.
 
         Returns
         -------
-        TIntScalar_co
+        Scalar
             The default value.
         """
         return self._cast_scalar_unchecked(0)
 
-    def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> TIntScalar_co:
+    def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> Scalar:
         """
         Read a JSON-serializable value as a NumPy int scalar.
 
@@ -198,7 +201,7 @@ def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> TIntScalar
 
         Returns
         -------
-        TIntScalar_co
+        Scalar
             The NumPy int scalar.
 
         Raises

src/zarr/core/dtype/npy/string.py

@@ -33,7 +33,7 @@
     endianness_to_numpy_str,
     get_endianness_from_numpy_dtype,
 )
-from zarr.core.dtype.wrapper import TDType_co, ZDType
+from zarr.core.dtype.wrapper import ZDType
 
 if TYPE_CHECKING:
     from zarr.core.common import JSON, ZarrFormat
@@ -453,7 +453,7 @@ class VariableLengthUTF8JSON_V2(DTypeConfig_V2[Literal["|O"], Literal["vlen-utf8
 # If NumPy 2 is installed, then VariableLengthUTF8 is defined with the NumPy variable length
 # string dtype as the native dtype. Otherwise, VariableLengthUTF8 is defined with the NumPy object
 # dtype as the native dtype.
-class UTF8Base(ZDType[TDType_co, str], HasObjectCodec):
+class UTF8Base[DType: TBaseDType](ZDType[DType, str], HasObjectCodec):
     """
     A base class for variable-length UTF-8 string data types.
 

src/zarr/core/dtype/npy/time.py

@@ -217,7 +217,10 @@ class DateTime64JSON_V2(DTypeConfig_V2[str, None]):
 
 
 @dataclass(frozen=True, kw_only=True, slots=True)
-class TimeDTypeBase(ZDType[BaseTimeDType_co, BaseTimeScalar_co], HasEndianness, HasItemSize):
+class TimeDTypeBase[
+    DType: np.dtypes.TimeDelta64DType | np.dtypes.DateTime64DType,
+    Scalar: np.timedelta64 | np.datetime64,
+](ZDType[DType, Scalar], HasEndianness, HasItemSize):
     """
     A base class for data types that represent time via the NumPy TimeDelta64 and DateTime64 data
     types.
@@ -275,7 +278,7 @@ def from_native_dtype(cls, dtype: TBaseDType) -> Self:
             f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}"
         )
 
-    def to_native_dtype(self) -> BaseTimeDType_co:
+    def to_native_dtype(self) -> DType:
         # Numpy does not allow creating datetime64 or timedelta64 via
         # np.dtypes.{dtype_name}()
         # so we use np.dtype with a formatted string.
@@ -285,7 +288,7 @@ def to_native_dtype(self) -> BaseTimeDType_co:
 
         Returns
         -------
-        BaseTimeDType_co
+        DType
             A NumPy data type object representing the time data type with
             the specified unit, scale factor, and byte order.
         """