Add chunkedduckarray to _typing (#8376)

Author: Illviljan

xarray/namedarray/_typing.py

@@ -102,10 +102,6 @@ def __array__(self, dtype: _DType, /) -> np.ndarray[Any, _DType]:
         ...
 
 
-# Corresponds to np.typing.NDArray:
-_Array = _array[Any, np.dtype[_ScalarType_co]]
-
-
 @runtime_checkable
 class _arrayfunction(
     _array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co]
@@ -147,10 +143,6 @@ def real(self) -> _arrayfunction[_ShapeType_co, Any]:
         ...
 
 
-# Corresponds to np.typing.NDArray:
-_ArrayFunction = _arrayfunction[Any, np.dtype[_ScalarType_co]]
-
-
 @runtime_checkable
 class _arrayapi(_array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co]):
     """
@@ -163,20 +155,15 @@ def __array_namespace__(self) -> ModuleType:
         ...
 
 
-# Corresponds to np.typing.NDArray:
-_ArrayAPI = _arrayapi[Any, np.dtype[_ScalarType_co]]
-
 # NamedArray can most likely use both __array_function__ and __array_namespace__:
 _arrayfunction_or_api = (_arrayfunction, _arrayapi)
-# _ArrayFunctionOrAPI = Union[
-#     _arrayfunction[_ShapeType_co, _DType_co], _arrayapi[_ShapeType_co, _DType_co]
-# ]
 
 duckarray = Union[
     _arrayfunction[_ShapeType_co, _DType_co], _arrayapi[_ShapeType_co, _DType_co]
 ]
+
+# Corresponds to np.typing.NDArray:
 DuckArray = _arrayfunction[Any, np.dtype[_ScalarType_co]]
-T_DuckArray = TypeVar("T_DuckArray", bound=_arrayfunction[Any, Any])
 
 
 @runtime_checkable
@@ -194,10 +181,6 @@ def chunks(self) -> _Chunks:
         ...
 
 
-# Corresponds to np.typing.NDArray:
-_ChunkedArray = _chunkedarray[Any, np.dtype[_ScalarType_co]]
-
-
 @runtime_checkable
 class _chunkedarrayfunction(
     _arrayfunction[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co]
@@ -213,10 +196,6 @@ def chunks(self) -> _Chunks:
         ...
 
 
-# Corresponds to np.typing.NDArray:
-_ChunkedArrayFunction = _chunkedarrayfunction[Any, np.dtype[_ScalarType_co]]
-
-
 @runtime_checkable
 class _chunkedarrayapi(
     _arrayapi[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co]
@@ -232,8 +211,12 @@ def chunks(self) -> _Chunks:
         ...
 
 
-# Corresponds to np.typing.NDArray:
-_ChunkedArrayAPI = _chunkedarrayapi[Any, np.dtype[_ScalarType_co]]
+# NamedArray can most likely use both __array_function__ and __array_namespace__:
+_chunkedarrayfunction_or_api = (_chunkedarrayfunction, _chunkedarrayapi)
+chunkedduckarray = Union[
+    _chunkedarrayfunction[_ShapeType_co, _DType_co],
+    _chunkedarrayapi[_ShapeType_co, _DType_co],
+]
 
 
 @runtime_checkable
@@ -250,10 +233,6 @@ def todense(self) -> NDArray[_ScalarType_co]:
         ...
 
 
-# Corresponds to np.typing.NDArray:
-_SparseArray = _sparsearray[Any, np.dtype[_ScalarType_co]]
-
-
 @runtime_checkable
 class _sparsearrayfunction(
     _arrayfunction[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co]
@@ -268,10 +247,6 @@ def todense(self) -> NDArray[_ScalarType_co]:
         ...
 
 
-# Corresponds to np.typing.NDArray:
-_SparseArrayFunction = _sparsearrayfunction[Any, np.dtype[_ScalarType_co]]
-
-
 @runtime_checkable
 class _sparsearrayapi(
     _arrayapi[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co]
@@ -286,22 +261,5 @@ def todense(self) -> NDArray[_ScalarType_co]:
         ...
 
 
-# Corresponds to np.typing.NDArray:
-_SparseArrayAPI = _sparsearrayapi[Any, np.dtype[_ScalarType_co]]
-
 # NamedArray can most likely use both __array_function__ and __array_namespace__:
 _sparsearrayfunction_or_api = (_sparsearrayfunction, _sparsearrayapi)
-_SparseArrayFunctionOrAPI = Union[
-    _SparseArrayFunction[np.generic], _SparseArrayAPI[np.generic]
-]
-
-
-# Temporary placeholder for indicating an array api compliant type.
-# hopefully in the future we can narrow this down more
-# T_DuckArray = TypeVar("T_DuckArray", bound=_ArrayFunctionOrAPI)
-
-# The chunked arrays like dask or cubed:
-_ChunkedArrayFunctionOrAPI = Union[
-    _ChunkedArrayFunction[np.generic], _ChunkedArrayAPI[np.generic]
-]
-T_ChunkedArray = TypeVar("T_ChunkedArray", bound=_ChunkedArrayFunctionOrAPI)