diff --git a/xarray/core/common.py b/xarray/core/common.py index 4c506b05325..bb3ccf44195 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -215,14 +215,13 @@ def __iter__(self: Any) -> Iterator[Any]: return self._iter() @overload - def get_axis_num(self, dim: str) -> int: ... # type: ignore [overload-overlap] + def get_axis_num( + self, dim: Hashable + ) -> int: ... # put this first to match a single str @overload def get_axis_num(self, dim: Iterable[Hashable]) -> tuple[int, ...]: ... - @overload - def get_axis_num(self, dim: Hashable) -> int: ... - def get_axis_num(self, dim: Hashable | Iterable[Hashable]) -> int | tuple[int, ...]: """Return axis number(s) corresponding to dimension(s) in this array. diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 149415847f3..2021c8a4de4 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -7254,7 +7254,7 @@ def _to_dataframe(self, ordered_dims: Mapping[Any, int]): extension_array_df = pd.DataFrame( {extension_array_column: extension_array}, index=pd.Index(index.array) - if isinstance(index, PandasExtensionArray) # type: ignore[redundant-expr] + if isinstance(index, PandasExtensionArray) else index, ) extension_array_df.index.name = self.variables[extension_array_column].dims[ diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index bb12704e55c..15461c14637 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -42,7 +42,7 @@ from xarray.core.indexes import Index from xarray.core.types import Self from xarray.core.variable import Variable - from xarray.namedarray._typing import _Shape, duckarray + from xarray.namedarray._typing import Shape, duckarray from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint BasicIndexerType = int | np.integer | slice @@ -730,7 +730,7 @@ def __init__(self, array: Any, key: ExplicitIndexer | None = None): self.array = as_indexable(array) self.key = key - shape: _Shape = () + shape: Shape = () for size, k in zip(self.array.shape, self.key.tuple, strict=True): if isinstance(k, slice): shape += (len(range(*k.indices(size))),) @@ -754,7 +754,7 @@ def _updated_key(self, new_key: ExplicitIndexer) -> BasicIndexer | OuterIndexer: return OuterIndexer(full_key_tuple) @property - def shape(self) -> _Shape: + def shape(self) -> Shape: return self._shape def get_duck_array(self): @@ -836,7 +836,7 @@ def __init__(self, array: duckarray[Any, Any], key: ExplicitIndexer): self.array = as_indexable(array) @property - def shape(self) -> _Shape: + def shape(self) -> Shape: return np.broadcast(*self.key.tuple).shape def get_duck_array(self): @@ -1025,7 +1025,7 @@ def as_indexable(array): def _outer_to_vectorized_indexer( - indexer: BasicIndexer | OuterIndexer, shape: _Shape + indexer: BasicIndexer | OuterIndexer, shape: Shape ) -> VectorizedIndexer: """Convert an OuterIndexer into a vectorized indexer. @@ -1061,7 +1061,7 @@ def _outer_to_vectorized_indexer( return VectorizedIndexer(tuple(new_key)) -def _outer_to_numpy_indexer(indexer: BasicIndexer | OuterIndexer, shape: _Shape): +def _outer_to_numpy_indexer(indexer: BasicIndexer | OuterIndexer, shape: Shape): """Convert an OuterIndexer into an indexer for NumPy. Parameters @@ -1085,7 +1085,7 @@ def _outer_to_numpy_indexer(indexer: BasicIndexer | OuterIndexer, shape: _Shape) return _outer_to_vectorized_indexer(indexer, shape).tuple -def _combine_indexers(old_key, shape: _Shape, new_key) -> VectorizedIndexer: +def _combine_indexers(old_key, shape: Shape, new_key) -> VectorizedIndexer: """Combine two indexers. Parameters @@ -1127,7 +1127,7 @@ class IndexingSupport(enum.Enum): def explicit_indexing_adapter( key: ExplicitIndexer, - shape: _Shape, + shape: Shape, indexing_support: IndexingSupport, raw_indexing_method: Callable[..., Any], ) -> Any: @@ -1163,7 +1163,7 @@ def explicit_indexing_adapter( async def async_explicit_indexing_adapter( key: ExplicitIndexer, - shape: _Shape, + shape: Shape, indexing_support: IndexingSupport, raw_indexing_method: Callable[..., Any], ) -> Any: @@ -1197,7 +1197,7 @@ def set_with_indexer(indexable, indexer: ExplicitIndexer, value: Any) -> None: def decompose_indexer( - indexer: ExplicitIndexer, shape: _Shape, indexing_support: IndexingSupport + indexer: ExplicitIndexer, shape: Shape, indexing_support: IndexingSupport ) -> tuple[ExplicitIndexer, ExplicitIndexer]: if isinstance(indexer, VectorizedIndexer): return _decompose_vectorized_indexer(indexer, shape, indexing_support) @@ -1236,7 +1236,7 @@ def _decompose_slice(key: slice, size: int) -> tuple[slice, slice]: def _decompose_vectorized_indexer( indexer: VectorizedIndexer, - shape: _Shape, + shape: Shape, indexing_support: IndexingSupport, ) -> tuple[ExplicitIndexer, ExplicitIndexer]: """ @@ -1318,7 +1318,7 @@ def _decompose_vectorized_indexer( def _decompose_outer_indexer( indexer: BasicIndexer | OuterIndexer, - shape: _Shape, + shape: Shape, indexing_support: IndexingSupport, ) -> tuple[ExplicitIndexer, ExplicitIndexer]: """ @@ -1500,7 +1500,7 @@ def _arrayize_outer_indexer(indexer: OuterIndexer, shape) -> OuterIndexer: def _arrayize_vectorized_indexer( - indexer: VectorizedIndexer, shape: _Shape + indexer: VectorizedIndexer, shape: Shape ) -> VectorizedIndexer: """Return an identical vindex but slices are replaced by arrays""" slices = [v for v in indexer.tuple if isinstance(v, slice)] @@ -1564,7 +1564,7 @@ def _masked_result_drop_slice(key, data: duckarray[Any, Any] | None = None): def create_mask( - indexer: ExplicitIndexer, shape: _Shape, data: duckarray[Any, Any] | None = None + indexer: ExplicitIndexer, shape: Shape, data: duckarray[Any, Any] | None = None ): """Create a mask for indexing with a fill-value. @@ -1975,7 +1975,7 @@ def get_duck_array(self) -> np.ndarray | PandasExtensionArray: return np.asarray(self) @property - def shape(self) -> _Shape: + def shape(self) -> Shape: return (len(self.array),) def _convert_scalar(self, item) -> np.ndarray: diff --git a/xarray/namedarray/_aggregations.py b/xarray/namedarray/_aggregations.py index c5726ef9251..ffb707eb1ea 100644 --- a/xarray/namedarray/_aggregations.py +++ b/xarray/namedarray/_aggregations.py @@ -5,19 +5,20 @@ from __future__ import annotations from collections.abc import Callable, Sequence -from typing import Any +from typing import Any, Generic from xarray.core import duck_array_ops -from xarray.core.types import Dims, Self +from xarray.core.types import Self +from xarray.namedarray._typing import DimsLike, DimType_co -class NamedArrayAggregations: +class NamedArrayAggregations(Generic[DimType_co]): __slots__ = () def reduce( self, func: Callable[..., Any], - dim: Dims = None, + dim: DimsLike[DimType_co] = None, *, axis: int | Sequence[int] | None = None, keepdims: bool = False, @@ -27,7 +28,7 @@ def reduce( def count( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, **kwargs: Any, ) -> Self: """ @@ -78,7 +79,7 @@ def count( def all( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, **kwargs: Any, ) -> Self: """ @@ -131,7 +132,7 @@ def all( def any( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, **kwargs: Any, ) -> Self: """ @@ -184,7 +185,7 @@ def any( def max( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, *, skipna: bool | None = None, **kwargs: Any, @@ -249,7 +250,7 @@ def max( def min( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, *, skipna: bool | None = None, **kwargs: Any, @@ -314,7 +315,7 @@ def min( def mean( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, *, skipna: bool | None = None, **kwargs: Any, @@ -379,7 +380,7 @@ def mean( def prod( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, *, skipna: bool | None = None, min_count: int | None = None, @@ -462,7 +463,7 @@ def prod( def sum( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, *, skipna: bool | None = None, min_count: int | None = None, @@ -545,7 +546,7 @@ def sum( def std( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, *, skipna: bool | None = None, ddof: int = 0, @@ -625,7 +626,7 @@ def std( def var( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, *, skipna: bool | None = None, ddof: int = 0, @@ -705,7 +706,7 @@ def var( def median( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, *, skipna: bool | None = None, **kwargs: Any, @@ -774,7 +775,7 @@ def median( def cumsum( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, *, skipna: bool | None = None, **kwargs: Any, @@ -848,7 +849,7 @@ def cumsum( def cumprod( self, - dim: Dims = None, + dim: DimsLike[DimType_co] = None, *, skipna: bool | None = None, **kwargs: Any, diff --git a/xarray/namedarray/_array_api.py b/xarray/namedarray/_array_api.py index 9cd064b110e..86f94ef3c76 100644 --- a/xarray/namedarray/_array_api.py +++ b/xarray/namedarray/_array_api.py @@ -1,28 +1,28 @@ from __future__ import annotations from types import ModuleType -from typing import Any +from typing import Any, overload import numpy as np from xarray.namedarray._typing import ( + Axes, + Axis, Default, - _arrayapi, - _Axes, - _Axis, + DimType, + DType, + ScalarType, + ShapeType, + SupportsImag, + SupportsReal, _default, - _Dim, - _DType, - _ScalarType, - _ShapeType, - _SupportsImag, - _SupportsReal, + arrayapi, ) from xarray.namedarray.core import NamedArray -def _get_data_namespace(x: NamedArray[Any, Any]) -> ModuleType: - if isinstance(x._data, _arrayapi): +def _get_data_namespace(x: NamedArray[Any, Any, Any]) -> ModuleType: + if isinstance(x._data, arrayapi): return x._data.__array_namespace__() return np @@ -32,8 +32,8 @@ def _get_data_namespace(x: NamedArray[Any, Any]) -> ModuleType: def astype( - x: NamedArray[_ShapeType, Any], dtype: _DType, /, *, copy: bool = True -) -> NamedArray[_ShapeType, _DType]: + x: NamedArray[ShapeType, Any, DimType], dtype: DType, /, *, copy: bool = True +) -> NamedArray[ShapeType, DType, DimType]: """ Copies an array to a specified data type irrespective of Type Promotion Rules rules. @@ -67,7 +67,7 @@ def astype( Size: 8B array([1, 2], dtype=int32) """ - if isinstance(x._data, _arrayapi): + if isinstance(x._data, arrayapi): xp = x._data.__array_namespace__() return x._new(data=xp.astype(x._data, dtype, copy=copy)) @@ -79,9 +79,9 @@ def astype( def imag( - x: NamedArray[_ShapeType, np.dtype[_SupportsImag[_ScalarType]]], # type: ignore[type-var] + x: NamedArray[ShapeType, np.dtype[SupportsImag[ScalarType]], DimType], # type: ignore[type-var] /, -) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]: +) -> NamedArray[ShapeType, np.dtype[ScalarType], DimType]: """ Returns the imaginary component of a complex number for each element x_i of the input array x. @@ -112,9 +112,9 @@ def imag( def real( - x: NamedArray[_ShapeType, np.dtype[_SupportsReal[_ScalarType]]], # type: ignore[type-var] + x: NamedArray[ShapeType, np.dtype[SupportsReal[ScalarType]], DimType], # type: ignore[type-var] /, -) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]: +) -> NamedArray[ShapeType, np.dtype[ScalarType], DimType]: """ Returns the real component of a complex number for each element x_i of the input array x. @@ -145,13 +145,35 @@ def real( # %% Manipulation functions + + +@overload +def expand_dims( + x: NamedArray[Any, DType, DimType], + /, + *, + dim: DimType, + axis: Axis = ..., +) -> NamedArray[Any, DType, DimType]: ... + + +@overload +def expand_dims( + x: NamedArray[Any, DType, DimType], + /, + *, + dim: Default = ..., + axis: Axis = ..., +) -> NamedArray[Any, DType, DimType | str]: ... + + def expand_dims( - x: NamedArray[Any, _DType], + x: NamedArray[Any, DType, DimType], /, *, - dim: _Dim | Default = _default, - axis: _Axis = 0, -) -> NamedArray[Any, _DType]: + dim: DimType | Default = _default, + axis: Axis = 0, +) -> NamedArray[Any, DType, DimType] | NamedArray[Any, DType, DimType | str]: """ Expands the shape of an array by inserting a new dimension of size one at the position specified by dims. @@ -184,15 +206,16 @@ def expand_dims( """ xp = _get_data_namespace(x) dims = x.dims - if dim is _default: - dim = f"dim_{len(dims)}" - d = list(dims) - d.insert(axis, dim) + actual_dim: DimType | str = f"dim_{len(dims)}" if dim is _default else dim + d: list[DimType | str] = list(dims) + d.insert(axis, actual_dim) out = x._new(dims=tuple(d), data=xp.expand_dims(x._data, axis=axis)) return out -def permute_dims(x: NamedArray[Any, _DType], axes: _Axes) -> NamedArray[Any, _DType]: +def permute_dims( + x: NamedArray[Any, DType, DimType], axes: Axes +) -> NamedArray[Any, DType, DimType]: """ Permutes the dimensions of an array. @@ -213,7 +236,7 @@ def permute_dims(x: NamedArray[Any, _DType], axes: _Axes) -> NamedArray[Any, _DT dims = x.dims new_dims = tuple(dims[i] for i in axes) - if isinstance(x._data, _arrayapi): + if isinstance(x._data, arrayapi): xp = _get_data_namespace(x) out = x._new(dims=new_dims, data=xp.permute_dims(x._data, axes)) else: diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 9610b96d4f9..6a5e857219e 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -37,75 +37,77 @@ class Default(Enum): # https://stackoverflow.com/questions/74633074/how-to-type-hint-a-generic-numpy-array _T_co = TypeVar("_T_co", covariant=True) -_dtype = np.dtype -_DType = TypeVar("_DType", bound=np.dtype[Any]) -_DType_co = TypeVar("_DType_co", covariant=True, bound=np.dtype[Any]) +dtype: TypeAlias = np.dtype # noqa: PYI042 +DType = TypeVar("DType", bound=np.dtype[Any]) +DType_co = TypeVar("DType_co", covariant=True, bound=np.dtype[Any]) # A subset of `npt.DTypeLike` that can be parametrized w.r.t. `np.generic` -_ScalarType = TypeVar("_ScalarType", bound=np.generic) -_ScalarType_co = TypeVar("_ScalarType_co", bound=np.generic, covariant=True) +ScalarType = TypeVar("ScalarType", bound=np.generic) +ScalarType_co = TypeVar("ScalarType_co", bound=np.generic, covariant=True) # A protocol for anything with the dtype attribute @runtime_checkable -class _SupportsDType(Protocol[_DType_co]): +class SupportsDType(Protocol[DType_co]): @property - def dtype(self) -> _DType_co: ... + def dtype(self) -> DType_co: ... -_DTypeLike = Union[ - np.dtype[_ScalarType], - type[_ScalarType], - _SupportsDType[np.dtype[_ScalarType]], +DTypeLike: TypeAlias = Union[ + np.dtype[ScalarType], + type[ScalarType], + SupportsDType[np.dtype[ScalarType]], ] # For unknown shapes Dask uses np.nan, array_api uses None: -_IntOrUnknown = int -_Shape = tuple[_IntOrUnknown, ...] -_ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]] -_ShapeType = TypeVar("_ShapeType", bound=Any) -_ShapeType_co = TypeVar("_ShapeType_co", bound=Any, covariant=True) - -_Axis = int -_Axes = tuple[_Axis, ...] -_AxisLike = Union[_Axis, _Axes] - -_Chunks = tuple[_Shape, ...] -_NormalizedChunks = tuple[tuple[int, ...], ...] +IntOrUnknown: TypeAlias = int +Shape: TypeAlias = tuple[IntOrUnknown, ...] +ShapeLike: TypeAlias = Union[SupportsIndex, Sequence[SupportsIndex]] +ShapeType = TypeVar("ShapeType", bound=Any) +ShapeType_co = TypeVar("ShapeType_co", bound=Any, covariant=True) + + +Axis: TypeAlias = int +Axes: TypeAlias = tuple[Axis, ...] +AxisLike = Union[Axis, Axes] + +Chunks: TypeAlias = tuple[Shape, ...] +NormalizedChunks: TypeAlias = tuple[tuple[int, ...], ...] # FYI in some cases we don't allow `None`, which this doesn't take account of. # # FYI the `str` is for a size string, e.g. "16MB", supported by dask. T_ChunkDim: TypeAlias = str | int | Literal["auto"] | tuple[int, ...] | None # noqa: PYI051 # We allow the tuple form of this (though arguably we could transition to named dims only) T_Chunks: TypeAlias = T_ChunkDim | Mapping[Any, T_ChunkDim] -_Dim = Hashable -_Dims = tuple[_Dim, ...] - -_DimsLike = Union[str, Iterable[_Dim]] +DimType = TypeVar("DimType", bound=Hashable) +DimType_co = TypeVar("DimType_co", bound=Hashable, covariant=True) +DimsLike: TypeAlias = Union[ + Iterable[DimType_co], None, EllipsisType +] # single str is also allowed, but luckily str = Iterable[str] # https://data-apis.org/array-api/latest/API_specification/indexing.html # TODO: np.array_api was bugged and didn't allow (None,), but should! # https://github.com/numpy/numpy/pull/25022 # https://github.com/data-apis/array-api/pull/674 -_IndexKey = Union[int, slice, EllipsisType] -_IndexKeys = tuple[_IndexKey, ...] # tuple[Union[_IndexKey, None], ...] -_IndexKeyLike = Union[_IndexKey, _IndexKeys] +IndexKey: TypeAlias = Union[int, slice, EllipsisType] +IndexKeys: TypeAlias = tuple[IndexKey, ...] # tuple[Union[_IndexKey, None], ...] +IndexKeyLike: TypeAlias = Union[IndexKey, IndexKeys] -_AttrsLike = Union[Mapping[Any, Any], None] +AttrsLike: TypeAlias = Union[Mapping[Any, Any], None] -class _SupportsReal(Protocol[_T_co]): +class SupportsReal(Protocol[_T_co]): @property def real(self) -> _T_co: ... -class _SupportsImag(Protocol[_T_co]): +class SupportsImag(Protocol[_T_co]): @property def imag(self) -> _T_co: ... @runtime_checkable -class _array(Protocol[_ShapeType_co, _DType_co]): +class array(Protocol[ShapeType_co, DType_co]): """ Minimal duck array named array uses. @@ -113,16 +115,14 @@ class _array(Protocol[_ShapeType_co, _DType_co]): """ @property - def shape(self) -> _Shape: ... + def shape(self) -> Shape: ... @property - def dtype(self) -> _DType_co: ... + def dtype(self) -> DType_co: ... @runtime_checkable -class _arrayfunction( - _array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] -): +class arrayfunction(array[ShapeType_co, DType_co], Protocol[ShapeType_co, DType_co]): """ Duck array supporting NEP 18. @@ -131,34 +131,33 @@ class _arrayfunction( @overload def __getitem__( - self, key: _arrayfunction[Any, Any] | tuple[_arrayfunction[Any, Any], ...], / - ) -> _arrayfunction[Any, _DType_co]: ... + self, key: arrayfunction[Any, Any] | tuple[arrayfunction[Any, Any], ...], / + ) -> arrayfunction[Any, DType_co]: ... @overload - def __getitem__(self, key: _IndexKeyLike, /) -> Any: ... + def __getitem__(self, key: IndexKeyLike, /) -> Any: ... def __getitem__( self, key: ( - _IndexKeyLike - | _arrayfunction[Any, Any] - | tuple[_arrayfunction[Any, Any], ...] + IndexKeyLike | arrayfunction[Any, Any] | tuple[arrayfunction[Any, Any], ...] ), /, - ) -> _arrayfunction[Any, _DType_co] | Any: ... + ) -> arrayfunction[Any, DType_co] | Any: ... @overload def __array__( self, dtype: None = ..., /, *, copy: bool | None = ... - ) -> np.ndarray[Any, _DType_co]: ... + ) -> np.ndarray[Any, DType_co]: ... + @overload def __array__( - self, dtype: _DType, /, *, copy: bool | None = ... - ) -> np.ndarray[Any, _DType]: ... + self, dtype: DType, /, *, copy: bool | None = ... + ) -> np.ndarray[Any, DType]: ... def __array__( - self, dtype: _DType | None = ..., /, *, copy: bool | None = ... - ) -> np.ndarray[Any, _DType] | np.ndarray[Any, _DType_co]: ... + self, dtype: DType | None = ..., /, *, copy: bool | None = ... + ) -> np.ndarray[Any, DType] | np.ndarray[Any, DType_co]: ... # TODO: Should return the same subclass but with a new dtype generic. # https://github.com/python/typing/issues/548 @@ -181,14 +180,14 @@ def __array_function__( ) -> Any: ... @property - def imag(self) -> _arrayfunction[_ShapeType_co, Any]: ... + def imag(self) -> arrayfunction[ShapeType_co, Any]: ... @property - def real(self) -> _arrayfunction[_ShapeType_co, Any]: ... + def real(self) -> arrayfunction[ShapeType_co, Any]: ... @runtime_checkable -class _arrayapi(_array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co]): +class arrayapi(array[ShapeType_co, DType_co], Protocol[ShapeType_co, DType_co]): """ Duck array supporting NEP 47. @@ -198,29 +197,27 @@ class _arrayapi(_array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType def __getitem__( self, key: ( - _IndexKeyLike | Any + IndexKeyLike | Any ), # TODO: Any should be _arrayapi[Any, _dtype[np.integer]] /, - ) -> _arrayapi[Any, Any]: ... + ) -> arrayapi[Any, Any]: ... def __array_namespace__(self) -> ModuleType: ... # NamedArray can most likely use both __array_function__ and __array_namespace__: -_arrayfunction_or_api = (_arrayfunction, _arrayapi) +_arrayfunction_or_api = (arrayfunction, arrayapi) -duckarray = Union[ - _arrayfunction[_ShapeType_co, _DType_co], _arrayapi[_ShapeType_co, _DType_co] +duckarray: TypeAlias = Union[ # noqa: PYI042 + arrayfunction[ShapeType_co, DType_co], arrayapi[ShapeType_co, DType_co] ] # Corresponds to np.typing.NDArray: -DuckArray = _arrayfunction[Any, np.dtype[_ScalarType_co]] +DuckArray: TypeAlias = arrayfunction[Any, np.dtype[ScalarType_co]] @runtime_checkable -class _chunkedarray( - _array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] -): +class chunkedarray(array[ShapeType_co, DType_co], Protocol[ShapeType_co, DType_co]): """ Minimal chunked duck array. @@ -228,12 +225,12 @@ class _chunkedarray( """ @property - def chunks(self) -> _Chunks: ... + def chunks(self) -> Chunks: ... @runtime_checkable -class _chunkedarrayfunction( - _arrayfunction[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] +class chunkedarrayfunction( + arrayfunction[ShapeType_co, DType_co], Protocol[ShapeType_co, DType_co] ): """ Chunked duck array supporting NEP 18. @@ -242,12 +239,12 @@ class _chunkedarrayfunction( """ @property - def chunks(self) -> _Chunks: ... + def chunks(self) -> Chunks: ... @runtime_checkable -class _chunkedarrayapi( - _arrayapi[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] +class chunkedarrayapi( + arrayapi[ShapeType_co, DType_co], Protocol[ShapeType_co, DType_co] ): """ Chunked duck array supporting NEP 47. @@ -256,33 +253,31 @@ class _chunkedarrayapi( """ @property - def chunks(self) -> _Chunks: ... + def chunks(self) -> Chunks: ... # 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], +_chunkedarrayfunction_or_api = (chunkedarrayfunction, chunkedarrayapi) +chunkedduckarray: TypeAlias = Union[ # noqa: PYI042 + chunkedarrayfunction[ShapeType_co, DType_co], + chunkedarrayapi[ShapeType_co, DType_co], ] @runtime_checkable -class _sparsearray( - _array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] -): +class sparsearray(array[ShapeType_co, DType_co], Protocol[ShapeType_co, DType_co]): """ Minimal sparse duck array. Corresponds to np.ndarray. """ - def todense(self) -> np.ndarray[Any, _DType_co]: ... + def todense(self) -> np.ndarray[Any, DType_co]: ... @runtime_checkable -class _sparsearrayfunction( - _arrayfunction[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] +class sparsearrayfunction( + arrayfunction[ShapeType_co, DType_co], Protocol[ShapeType_co, DType_co] ): """ Sparse duck array supporting NEP 18. @@ -290,12 +285,12 @@ class _sparsearrayfunction( Corresponds to np.ndarray. """ - def todense(self) -> np.ndarray[Any, _DType_co]: ... + def todense(self) -> np.ndarray[Any, DType_co]: ... @runtime_checkable -class _sparsearrayapi( - _arrayapi[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] +class sparsearrayapi( + arrayapi[ShapeType_co, DType_co], Protocol[ShapeType_co, DType_co] ): """ Sparse duck array supporting NEP 47. @@ -303,15 +298,15 @@ class _sparsearrayapi( Corresponds to np.ndarray. """ - def todense(self) -> np.ndarray[Any, _DType_co]: ... + def todense(self) -> np.ndarray[Any, DType_co]: ... # NamedArray can most likely use both __array_function__ and __array_namespace__: -_sparsearrayfunction_or_api = (_sparsearrayfunction, _sparsearrayapi) -sparseduckarray = Union[ - _sparsearrayfunction[_ShapeType_co, _DType_co], - _sparsearrayapi[_ShapeType_co, _DType_co], +_sparsearrayfunction_or_api = (sparsearrayfunction, sparsearrayapi) +sparseduckarray: TypeAlias = Union[ # noqa: PYI042 + sparsearrayfunction[ShapeType_co, DType_co], + sparsearrayapi[ShapeType_co, DType_co], ] -ErrorOptions = Literal["raise", "ignore"] -ErrorOptionsWithWarn = Literal["raise", "warn", "ignore"] +ErrorHandling: TypeAlias = Literal["raise", "ignore"] +ErrorHandlingWithWarn: TypeAlias = Literal["raise", "warn", "ignore"] diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py index 23d55ee0a11..6b42b7f85a3 100644 --- a/xarray/namedarray/core.py +++ b/xarray/namedarray/core.py @@ -3,7 +3,7 @@ import copy import math import warnings -from collections.abc import Callable, Hashable, Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from itertools import starmap from types import EllipsisType from typing import ( @@ -27,18 +27,19 @@ ) from xarray.namedarray._aggregations import NamedArrayAggregations from xarray.namedarray._typing import ( - ErrorOptionsWithWarn, - _arrayapi, + DimType, + DimType_co, + DType_co, + ErrorHandlingWithWarn, + ScalarType_co, + ShapeType_co, + SupportsImag, + SupportsReal, _arrayfunction_or_api, - _chunkedarray, _default, - _dtype, - _DType_co, - _ScalarType_co, - _ShapeType_co, _sparsearrayfunction_or_api, - _SupportsImag, - _SupportsReal, + arrayapi, + chunkedarray, ) from xarray.namedarray.parallelcompat import guess_chunkmanager from xarray.namedarray.pycompat import to_numpy @@ -53,19 +54,17 @@ if TYPE_CHECKING: from numpy.typing import ArrayLike, NDArray - from xarray.core.types import Dims, T_Chunks + from xarray.core.types import T_Chunks from xarray.namedarray._typing import ( + AttrsLike, + Chunks, Default, - _AttrsLike, - _Chunks, - _Dim, - _Dims, - _DimsLike, - _DType, - _IntOrUnknown, - _ScalarType, - _Shape, - _ShapeType, + DimsLike, + DType, + IntOrUnknown, + ScalarType, + Shape, + ShapeType, duckarray, ) from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint @@ -87,36 +86,59 @@ from typing import Self - T_NamedArray = TypeVar("T_NamedArray", bound="_NamedArray[Any]") + T_NamedArray = TypeVar("T_NamedArray", bound="_NamedArray[Any, Any]") T_NamedArrayInteger = TypeVar( - "T_NamedArrayInteger", bound="_NamedArray[np.integer[Any]]" + "T_NamedArrayInteger", bound="_NamedArray[np.integer[Any], Any]" ) @overload def _new( - x: NamedArray[Any, _DType_co], - dims: _DimsLike | Default = ..., - data: duckarray[_ShapeType, _DType] = ..., - attrs: _AttrsLike | Default = ..., -) -> NamedArray[_ShapeType, _DType]: ... + x: NamedArray[Any, DType_co, DimType_co], + dims: Default = ..., + data: duckarray[ShapeType, DType] = ..., + attrs: AttrsLike | Default = ..., +) -> NamedArray[ShapeType, DType, DimType_co]: ... @overload def _new( - x: NamedArray[_ShapeType_co, _DType_co], - dims: _DimsLike | Default = ..., + x: NamedArray[Any, DType_co, DimType_co], + dims: Iterable[DimType] = ..., + data: duckarray[ShapeType, DType] = ..., + attrs: AttrsLike | Default = ..., +) -> NamedArray[ShapeType, DType, DimType]: ... + + +@overload +def _new( + x: NamedArray[ShapeType_co, DType_co, DimType_co], + dims: Default = ..., + data: Default = ..., + attrs: AttrsLike | Default = ..., +) -> NamedArray[ShapeType_co, DType_co, DimType_co]: ... + + +@overload +def _new( + x: NamedArray[ShapeType_co, DType_co, DimType_co], + dims: Iterable[DimType] = ..., data: Default = ..., - attrs: _AttrsLike | Default = ..., -) -> NamedArray[_ShapeType_co, _DType_co]: ... + attrs: AttrsLike | Default = ..., +) -> NamedArray[ShapeType_co, DType_co, DimType]: ... def _new( - x: NamedArray[Any, _DType_co], - dims: _DimsLike | Default = _default, - data: duckarray[_ShapeType, _DType] | Default = _default, - attrs: _AttrsLike | Default = _default, -) -> NamedArray[_ShapeType, _DType] | NamedArray[Any, _DType_co]: + x: NamedArray[Any, DType_co, DimType_co], + dims: Iterable[DimType] | Default = _default, + data: duckarray[ShapeType, DType] | Default = _default, + attrs: AttrsLike | Default = _default, +) -> ( + NamedArray[ShapeType, DType, DimType] + | NamedArray[ShapeType, DType, DimType_co] + | NamedArray[Any, DType_co, DimType_co] + | NamedArray[Any, DType_co, DimType] +): """ Create a new array with new typing information. @@ -137,41 +159,33 @@ def _new( Will copy the attrs from x by default. """ dims_ = copy.copy(x._dims) if dims is _default else dims + attrs_ = copy.copy(x._attrs) if attrs is _default else attrs + data_ = copy.copy(x._data) if data is _default else data - attrs_: Mapping[Any, Any] | None - if attrs is _default: - attrs_ = None if x._attrs is None else x._attrs.copy() - else: - attrs_ = attrs - - if data is _default: - return type(x)(dims_, copy.copy(x._data), attrs_) - else: - cls_ = cast("type[NamedArray[_ShapeType, _DType]]", type(x)) - return cls_(dims_, data, attrs_) + return type(x)(dims_, data_, attrs_) # type: ignore[arg-type] @overload def from_array( - dims: _DimsLike, - data: duckarray[_ShapeType, _DType], - attrs: _AttrsLike = ..., -) -> NamedArray[_ShapeType, _DType]: ... + dims: Iterable[DimType_co], + data: duckarray[ShapeType, DType], + attrs: AttrsLike = ..., +) -> NamedArray[ShapeType, DType, DimType_co]: ... @overload def from_array( - dims: _DimsLike, + dims: Iterable[DimType_co], data: ArrayLike, - attrs: _AttrsLike = ..., -) -> NamedArray[Any, Any]: ... + attrs: AttrsLike = ..., +) -> NamedArray[Any, Any, DimType_co]: ... def from_array( - dims: _DimsLike, - data: duckarray[_ShapeType, _DType] | ArrayLike, - attrs: _AttrsLike = None, -) -> NamedArray[_ShapeType, _DType] | NamedArray[Any, Any]: + dims: Iterable[DimType_co], + data: duckarray[ShapeType, DType] | ArrayLike, + attrs: AttrsLike = None, +) -> NamedArray[ShapeType, DType, DimType_co] | NamedArray[Any, Any, DimType_co]: """ Create a Named array from an array-like object. @@ -212,7 +226,9 @@ def from_array( return NamedArray(dims, np.asarray(data), attrs) -class NamedArray(NamedArrayAggregations, Generic[_ShapeType_co, _DType_co]): +class NamedArray( + NamedArrayAggregations[DimType_co], Generic[ShapeType_co, DType_co, DimType_co] +): """ A wrapper around duck arrays with named dimensions and attributes which describe a single Array. @@ -247,16 +263,16 @@ class NamedArray(NamedArrayAggregations, Generic[_ShapeType_co, _DType_co]): __slots__ = ("_attrs", "_data", "_dims") - _data: duckarray[Any, _DType_co] - _dims: _Dims + _data: duckarray[Any, DType_co] + _dims: tuple[DimType_co, ...] _attrs: dict[Any, Any] | None def __init__( self, - dims: _DimsLike, - data: duckarray[Any, _DType_co], - attrs: _AttrsLike = None, - ): + dims: Iterable[DimType_co], # str = Iterable[str] + data: duckarray[Any, DType_co], + attrs: AttrsLike = None, + ) -> None: self._data = data self._dims = self._parse_dimensions(dims) self._attrs = dict(attrs) if attrs else None @@ -273,25 +289,46 @@ def __init_subclass__(cls, **kwargs: Any) -> None: @overload def _new( self, - dims: _DimsLike | Default = ..., - data: duckarray[_ShapeType, _DType] = ..., - attrs: _AttrsLike | Default = ..., - ) -> NamedArray[_ShapeType, _DType]: ... + dims: Default = ..., + data: duckarray[ShapeType, DType] = ..., + attrs: AttrsLike | Default = ..., + ) -> NamedArray[ShapeType, DType, DimType_co]: ... + + @overload + def _new( + self, + dims: Iterable[DimType] = ..., + data: duckarray[ShapeType, DType] = ..., + attrs: AttrsLike | Default = ..., + ) -> NamedArray[ShapeType, DType, DimType]: ... @overload def _new( self, - dims: _DimsLike | Default = ..., + dims: Default = ..., data: Default = ..., - attrs: _AttrsLike | Default = ..., - ) -> NamedArray[_ShapeType_co, _DType_co]: ... + attrs: AttrsLike | Default = ..., + ) -> NamedArray[ShapeType_co, DType_co, DimType_co]: ... + @overload def _new( self, - dims: _DimsLike | Default = _default, - data: duckarray[Any, _DType] | Default = _default, - attrs: _AttrsLike | Default = _default, - ) -> NamedArray[_ShapeType, _DType] | NamedArray[_ShapeType_co, _DType_co]: + dims: Iterable[DimType] = ..., + data: Default = ..., + attrs: AttrsLike | Default = ..., + ) -> NamedArray[ShapeType_co, DType_co, DimType]: ... + + def _new( + self, + dims: Iterable[DimType] | Default = _default, + data: duckarray[Any, DType] | Default = _default, + attrs: AttrsLike | Default = _default, + ) -> ( + NamedArray[ShapeType, DType, DimType_co] + | NamedArray[ShapeType, DType, DimType] + | NamedArray[ShapeType_co, DType_co, DimType_co] + | NamedArray[ShapeType_co, DType_co, DimType] + ): """ Create a new array with new typing information. @@ -317,9 +354,9 @@ def _new( def _replace( self, - dims: _DimsLike | Default = _default, - data: duckarray[_ShapeType_co, _DType_co] | Default = _default, - attrs: _AttrsLike | Default = _default, + dims: Iterable[DimType] | Default = _default, + data: duckarray[ShapeType_co, DType_co] | Default = _default, + attrs: AttrsLike | Default = _default, ) -> Self: """ Create a new array with the same typing information. @@ -346,7 +383,7 @@ def _replace( def _copy( self, deep: bool = True, - data: duckarray[_ShapeType_co, _DType_co] | None = None, + data: duckarray[ShapeType_co, DType_co] | None = None, memo: dict[int, Any] | None = None, ) -> Self: if data is None: @@ -372,7 +409,7 @@ def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Self: def copy( self, deep: bool = True, - data: duckarray[_ShapeType_co, _DType_co] | None = None, + data: duckarray[ShapeType_co, DType_co] | None = None, ) -> Self: """Returns a copy of this object. @@ -413,7 +450,7 @@ def ndim(self) -> int: return len(self.shape) @property - def size(self) -> _IntOrUnknown: + def size(self) -> IntOrUnknown: """ Number of elements in the array. @@ -425,14 +462,14 @@ def size(self) -> _IntOrUnknown: """ return math.prod(self.shape) - def __len__(self) -> _IntOrUnknown: + def __len__(self) -> IntOrUnknown: try: return self.shape[0] except Exception as exc: raise TypeError("len() of unsized object") from exc @property - def dtype(self) -> _DType_co: + def dtype(self) -> DType_co: """ Data-type of the array’s elements. @@ -444,7 +481,7 @@ def dtype(self) -> _DType_co: return self._data.dtype @property - def shape(self) -> _Shape: + def shape(self) -> Shape: """ Get the shape of the array. @@ -460,7 +497,7 @@ def shape(self) -> _Shape: return self._data.shape @property - def nbytes(self) -> _IntOrUnknown: + def nbytes(self) -> IntOrUnknown: """ Total bytes consumed by the elements of the data array. @@ -474,7 +511,7 @@ def nbytes(self) -> _IntOrUnknown: if hasattr(self.dtype, "itemsize"): itemsize = self.dtype.itemsize - elif isinstance(self._data, _arrayapi): + elif isinstance(self._data, arrayapi): xp = _get_data_namespace(self) if xp.isdtype(self.dtype, "bool"): @@ -491,32 +528,34 @@ def nbytes(self) -> _IntOrUnknown: return self.size * itemsize @property - def dims(self) -> _Dims: + def dims(self) -> tuple[DimType_co, ...]: """Tuple of dimension names with which this NamedArray is associated.""" return self._dims @dims.setter - def dims(self, value: _DimsLike) -> None: + def dims(self, value: Iterable[DimType_co]) -> None: self._dims = self._parse_dimensions(value) - def _parse_dimensions(self, dims: _DimsLike) -> _Dims: - dims = (dims,) if isinstance(dims, str) else tuple(dims) - if len(dims) != self.ndim: + def _parse_dimensions(self, dims: Iterable[DimType_co]) -> tuple[DimType_co, ...]: + dims_tuple = cast( + tuple[DimType_co, ...], (dims,) if isinstance(dims, str) else tuple(dims) + ) + if len(dims_tuple) != self.ndim: raise ValueError( f"dimensions {dims} must have the same length as the " f"number of data dimensions, ndim={self.ndim}" ) - if len(set(dims)) < len(dims): - repeated_dims = {d for d in dims if dims.count(d) > 1} + if len(set(dims_tuple)) < len(dims_tuple): + repeated_dims = {d for d in dims_tuple if dims_tuple.count(d) > 1} warnings.warn( - f"Duplicate dimension names present: dimensions {repeated_dims} appear more than once in dims={dims}. " + f"Duplicate dimension names present: dimensions {repeated_dims} appear more than once in dims={dims_tuple}. " "We do not yet support duplicate dimension names, but we do allow initial construction of the object. " "We recommend you rename the dims immediately to become distinct, as most xarray functionality is likely to fail silently if you do not. " "To rename the dimensions you will need to set the ``.dims`` attribute of each variable, ``e.g. var.dims=('x0', 'x1')``.", UserWarning, stacklevel=2, ) - return dims + return dims_tuple @property def attrs(self) -> dict[Any, Any]: @@ -529,7 +568,7 @@ def attrs(self) -> dict[Any, Any]: def attrs(self, value: Mapping[Any, Any]) -> None: self._attrs = dict(value) if value else None - def _check_shape(self, new_data: duckarray[Any, _DType_co]) -> None: + def _check_shape(self, new_data: duckarray[Any, DType_co]) -> None: if new_data.shape != self.shape: raise ValueError( f"replacement data must match the {self.__class__.__name__}'s shape. " @@ -537,7 +576,7 @@ def _check_shape(self, new_data: duckarray[Any, _DType_co]) -> None: ) @property - def data(self) -> duckarray[Any, _DType_co]: + def data(self) -> duckarray[Any, DType_co]: """ The NamedArray's data as an array. The underlying array type (e.g. dask, sparse, pint) is preserved. @@ -547,14 +586,14 @@ def data(self) -> duckarray[Any, _DType_co]: return self._data @data.setter - def data(self, data: duckarray[Any, _DType_co]) -> None: + def data(self, data: duckarray[Any, DType_co]) -> None: self._check_shape(data) self._data = data @property def imag( - self: NamedArray[_ShapeType, np.dtype[_SupportsImag[_ScalarType]]], # type: ignore[type-var] - ) -> NamedArray[_ShapeType, _dtype[_ScalarType]]: + self: NamedArray[ShapeType, np.dtype[SupportsImag[ScalarType]], DimType_co], # type: ignore[type-var] + ) -> NamedArray[ShapeType, np.dtype[ScalarType], DimType_co]: """ The imaginary part of the array. @@ -562,7 +601,7 @@ def imag( -------- numpy.ndarray.imag """ - if isinstance(self._data, _arrayapi): + if isinstance(self._data, arrayapi): from xarray.namedarray._array_api import imag return imag(self) @@ -571,8 +610,8 @@ def imag( @property def real( - self: NamedArray[_ShapeType, np.dtype[_SupportsReal[_ScalarType]]], # type: ignore[type-var] - ) -> NamedArray[_ShapeType, _dtype[_ScalarType]]: + self: NamedArray[ShapeType, np.dtype[SupportsReal[ScalarType]], DimType_co], # type: ignore[type-var] + ) -> NamedArray[ShapeType, np.dtype[ScalarType], DimType_co]: """ The real part of the array. @@ -580,7 +619,7 @@ def real( -------- numpy.ndarray.real """ - if isinstance(self._data, _arrayapi): + if isinstance(self._data, arrayapi): from xarray.namedarray._array_api import real return real(self) @@ -667,15 +706,14 @@ def _dask_finalize( return type(self)(self._dims, data, attrs=self._attrs) @overload - def get_axis_num(self, dim: str) -> int: ... # type: ignore [overload-overlap] - - @overload - def get_axis_num(self, dim: Iterable[Hashable]) -> tuple[int, ...]: ... + def get_axis_num(self, dim: DimType_co) -> int: ... # type: ignore[misc] # put this first to match a single str @overload - def get_axis_num(self, dim: Hashable) -> int: ... + def get_axis_num(self, dim: Iterable[DimType_co]) -> tuple[int, ...]: ... - def get_axis_num(self, dim: Hashable | Iterable[Hashable]) -> int | tuple[int, ...]: + def get_axis_num( + self, dim: DimType_co | Iterable[DimType_co] + ) -> int | tuple[int, ...]: """Return axis number(s) corresponding to dimension(s) in this array. Parameters @@ -693,17 +731,17 @@ def get_axis_num(self, dim: Hashable | Iterable[Hashable]) -> int | tuple[int, . else: return self._get_axis_num(dim) - def _get_axis_num(self: Any, dim: Hashable) -> int: + def _get_axis_num(self, dim: DimType_co) -> int: # type: ignore[misc] _raise_if_any_duplicate_dimensions(self.dims) try: - return self.dims.index(dim) # type: ignore[no-any-return] + return self.dims.index(dim) except ValueError as err: raise ValueError( f"{dim!r} not found in array dimensions {self.dims!r}" ) from err @property - def chunks(self) -> _Chunks | None: + def chunks(self) -> Chunks | None: """ Tuple of block lengths for this NamedArray's data, in order of dimensions, or None if the underlying data is not a dask array. @@ -715,7 +753,7 @@ def chunks(self) -> _Chunks | None: xarray.unify_chunks """ data = self._data - if isinstance(data, _chunkedarray): + if isinstance(data, chunkedarray): return data.chunks else: return None @@ -723,7 +761,7 @@ def chunks(self) -> _Chunks | None: @property def chunksizes( self, - ) -> Mapping[_Dim, _Shape]: + ) -> Mapping[DimType_co, Shape]: """ Mapping from dimension names to block lengths for this NamedArray's data. @@ -741,13 +779,13 @@ def chunksizes( xarray.unify_chunks """ data = self._data - if isinstance(data, _chunkedarray): + if isinstance(data, chunkedarray): return dict(zip(self.dims, data.chunks, strict=True)) else: return {} @property - def sizes(self) -> dict[_Dim, _IntOrUnknown]: + def sizes(self) -> dict[DimType_co, IntOrUnknown]: """Ordered mapping from dimension names to lengths.""" return dict(zip(self.dims, self.shape, strict=True)) @@ -864,11 +902,11 @@ def as_numpy(self) -> Self: def reduce( self, func: Callable[..., Any], - dim: Dims = None, + dim: DimsLike[DimType_co] = None, axis: int | Sequence[int] | None = None, keepdims: bool = False, **kwargs: Any, - ) -> NamedArray[Any, Any]: + ) -> NamedArray[Any, Any, DimType_co]: """Reduce this array by applying `func` along some dimension(s). Parameters @@ -903,7 +941,7 @@ def reduce( raise ValueError("cannot supply both 'axis' and 'dim' arguments") if dim is not None: - axis = self.get_axis_num(dim) + axis = self.get_axis_num(cast(Iterable[DimType_co], dim)) with warnings.catch_warnings(): warnings.filterwarnings( @@ -968,7 +1006,7 @@ def _as_sparse( self, sparse_format: Literal["coo"] | Default = _default, fill_value: ArrayLike | Default = _default, - ) -> NamedArray[Any, _DType_co]: + ) -> NamedArray[Any, DType_co, DimType_co]: """ Use sparse-array as backend. """ @@ -992,26 +1030,26 @@ def _as_sparse( data = as_sparse(astype(self, dtype).data, fill_value=fill_value) return self._new(data=data) - def _to_dense(self) -> NamedArray[Any, _DType_co]: + def _to_dense(self) -> NamedArray[Any, DType_co, DimType_co]: """ Change backend from sparse to np.array. """ if isinstance(self._data, _sparsearrayfunction_or_api): - data_dense: np.ndarray[Any, _DType_co] = self._data.todense() + data_dense: np.ndarray[Any, DType_co] = self._data.todense() return self._new(data=data_dense) else: raise TypeError("self.data is not a sparse array") def permute_dims( self, - *dim: Iterable[_Dim] | EllipsisType, - missing_dims: ErrorOptionsWithWarn = "raise", - ) -> NamedArray[Any, _DType_co]: + *dim: DimType_co | EllipsisType, + missing_dims: ErrorHandlingWithWarn = "raise", + ) -> NamedArray[Any, DType_co, DimType_co]: """Return a new object with transposed dimensions. Parameters ---------- - *dim : Hashable, optional + *dim : Hashable, "...", optional By default, reverse the order of the dimensions. Otherwise, reorder the dimensions to this order. missing_dims : {"raise", "warn", "ignore"}, default: "raise" @@ -1038,7 +1076,7 @@ def permute_dims( if not dim: dims = self.dims[::-1] else: - dims = tuple(infix_dims(dim, self.dims, missing_dims)) # type: ignore[arg-type] + dims = tuple(infix_dims(dim, self.dims, missing_dims)) if len(dims) < 2 or dims == self.dims: # no need to transpose if only one dimension @@ -1051,7 +1089,7 @@ def permute_dims( return permute_dims(self, axes) @property - def T(self) -> NamedArray[Any, _DType_co]: + def T(self) -> NamedArray[Any, DType_co, DimType_co]: """Return a new object with transposed dimensions.""" if self.ndim != 2: raise ValueError( @@ -1061,8 +1099,8 @@ def T(self) -> NamedArray[Any, _DType_co]: return self.permute_dims() def broadcast_to( - self, dim: Mapping[_Dim, int] | None = None, **dim_kwargs: Any - ) -> NamedArray[Any, _DType_co]: + self, dim: Mapping[DimType_co, int] | None = None, **dim_kwargs: int + ) -> NamedArray[Any, DType_co, DimType_co]: """ Broadcast the NamedArray to a new shape. New dimensions are not allowed. @@ -1120,10 +1158,25 @@ def broadcast_to( data = duck_array_ops.broadcast_to(self._data, ordered_shape) # type: ignore[no-untyped-call] # TODO: use array-api-compat function return self._new(data=data, dims=ordered_dims) + @overload + def expand_dims( + self, + dim: DimType_co, # type: ignore[misc] + ) -> NamedArray[Any, DType_co, DimType_co]: ... + + @overload + def expand_dims( + self, + dim: Default = ..., + ) -> NamedArray[Any, DType_co, DimType_co | str]: ... + def expand_dims( self, - dim: _Dim | Default = _default, - ) -> NamedArray[Any, _DType_co]: + dim: DimType_co | Default = _default, + ) -> ( + NamedArray[Any, DType_co, DimType_co] + | NamedArray[Any, DType_co, DimType_co | str] + ): """ Expand the dimensions of the NamedArray. @@ -1159,11 +1212,11 @@ def expand_dims( return expand_dims(self, dim=dim) -_NamedArray = NamedArray[Any, np.dtype[_ScalarType_co]] +_NamedArray = NamedArray[Any, np.dtype[ScalarType_co], DimType_co] def _raise_if_any_duplicate_dimensions( - dims: _Dims, err_context: str = "This function" + dims: tuple[DimType_co, ...], err_context: str = "This function" ) -> None: if len(set(dims)) < len(dims): repeated_dims = {d for d in dims if dims.count(d) > 1} diff --git a/xarray/namedarray/daskmanager.py b/xarray/namedarray/daskmanager.py index eb01a150c18..e082e2bb0c7 100644 --- a/xarray/namedarray/daskmanager.py +++ b/xarray/namedarray/daskmanager.py @@ -11,9 +11,9 @@ if TYPE_CHECKING: from xarray.namedarray._typing import ( + DType_co, + NormalizedChunks, T_Chunks, - _DType_co, - _NormalizedChunks, duckarray, ) @@ -40,16 +40,16 @@ def __init__(self) -> None: def is_chunked_array(self, data: duckarray[Any, Any]) -> bool: return is_duck_dask_array(data) - def chunks(self, data: Any) -> _NormalizedChunks: + def chunks(self, data: Any) -> NormalizedChunks: return data.chunks # type: ignore[no-any-return] def normalize_chunks( self, - chunks: T_Chunks | _NormalizedChunks, + chunks: T_Chunks | NormalizedChunks, shape: tuple[int, ...] | None = None, limit: int | None = None, - dtype: _DType_co | None = None, - previous_chunks: _NormalizedChunks | None = None, + dtype: DType_co | None = None, + previous_chunks: NormalizedChunks | None = None, ) -> Any: """Called by open_dataset""" from dask.array.core import normalize_chunks @@ -63,7 +63,7 @@ def normalize_chunks( ) # type: ignore[no-untyped-call] def from_array( - self, data: Any, chunks: T_Chunks | _NormalizedChunks, **kwargs: Any + self, data: Any, chunks: T_Chunks | NormalizedChunks, **kwargs: Any ) -> DaskArray | Any: import dask.array as da @@ -79,7 +79,7 @@ def from_array( def compute( self, *data: Any, **kwargs: Any - ) -> tuple[np.ndarray[Any, _DType_co], ...]: + ) -> tuple[np.ndarray[Any, DType_co], ...]: from dask.array import compute return compute(*data, **kwargs) # type: ignore[no-untyped-call, no-any-return] @@ -102,7 +102,7 @@ def reduction( combine_func: Callable[..., Any] | None = None, aggregate_func: Callable[..., Any] | None = None, axis: int | Sequence[int] | None = None, - dtype: _DType_co | None = None, + dtype: DType_co | None = None, keepdims: bool = False, ) -> DaskArray | Any: from dask.array import reduction @@ -124,7 +124,7 @@ def scan( ident: float, arr: T_ChunkedArray, axis: int | None = None, - dtype: _DType_co | None = None, + dtype: DType_co | None = None, **kwargs: Any, ) -> DaskArray | Any: from dask.array.reductions import cumreduction @@ -147,11 +147,11 @@ def apply_gufunc( axes: Sequence[tuple[int, ...]] | None = None, axis: int | None = None, keepdims: bool = False, - output_dtypes: Sequence[_DType_co] | None = None, + output_dtypes: Sequence[DType_co] | None = None, output_sizes: dict[str, int] | None = None, vectorize: bool | None = None, allow_rechunk: bool = False, - meta: tuple[np.ndarray[Any, _DType_co], ...] | None = None, + meta: tuple[np.ndarray[Any, DType_co], ...] | None = None, **kwargs: Any, ) -> Any: from dask.array.gufunc import apply_gufunc @@ -175,7 +175,7 @@ def map_blocks( self, func: Callable[..., Any], *args: Any, - dtype: _DType_co | None = None, + dtype: DType_co | None = None, chunks: tuple[int, ...] | None = None, drop_axis: int | Sequence[int] | None = None, new_axis: int | Sequence[int] | None = None, @@ -202,12 +202,12 @@ def blockwise( # can't type this as mypy assumes args are all same type, but dask blockwise args alternate types name: str | None = None, token: Any | None = None, - dtype: _DType_co | None = None, + dtype: DType_co | None = None, adjust_chunks: dict[Any, Callable[..., Any]] | None = None, new_axes: dict[Any, int] | None = None, align_arrays: bool = True, concatenate: bool | None = None, - meta: tuple[np.ndarray[Any, _DType_co], ...] | None = None, + meta: tuple[np.ndarray[Any, DType_co], ...] | None = None, **kwargs: Any, ) -> DaskArray | Any: from dask.array import blockwise @@ -231,7 +231,7 @@ def unify_chunks( self, *args: Any, # can't type this as mypy assumes args are all same type, but dask unify_chunks args alternate types **kwargs: Any, - ) -> tuple[dict[str, _NormalizedChunks], list[DaskArray]]: + ) -> tuple[dict[str, NormalizedChunks], list[DaskArray]]: from dask.array.core import unify_chunks return unify_chunks(*args, **kwargs) # type: ignore[no-any-return, no-untyped-call] diff --git a/xarray/namedarray/parallelcompat.py b/xarray/namedarray/parallelcompat.py index 8a68f5e9562..a1f90a50ecc 100644 --- a/xarray/namedarray/parallelcompat.py +++ b/xarray/namedarray/parallelcompat.py @@ -20,12 +20,12 @@ if TYPE_CHECKING: from xarray.namedarray._typing import ( + Chunks, + DType, + DType_co, + NormalizedChunks, + ShapeType, T_Chunks, - _Chunks, - _DType, - _DType_co, - _NormalizedChunks, - _ShapeType, duckarray, ) @@ -37,11 +37,11 @@ def rechunk(self, chunks: Any, **kwargs: Any) -> Any: ... def dtype(self) -> np.dtype[Any]: ... @property - def chunks(self) -> _NormalizedChunks: ... + def chunks(self) -> NormalizedChunks: ... def compute( self, *data: Any, **kwargs: Any - ) -> tuple[np.ndarray[Any, _DType_co], ...]: ... + ) -> tuple[np.ndarray[Any, DType_co], ...]: ... T_ChunkedArray = TypeVar("T_ChunkedArray", bound=ChunkedArrayMixinProtocol) @@ -238,7 +238,7 @@ def is_chunked_array(self, data: duckarray[Any, Any]) -> bool: return isinstance(data, self.array_cls) @abstractmethod - def chunks(self, data: T_ChunkedArray) -> _NormalizedChunks: + def chunks(self, data: T_ChunkedArray) -> NormalizedChunks: """ Return the current chunks of the given array. @@ -264,12 +264,12 @@ def chunks(self, data: T_ChunkedArray) -> _NormalizedChunks: @abstractmethod def normalize_chunks( self, - chunks: _Chunks | _NormalizedChunks, - shape: _ShapeType | None = None, + chunks: Chunks | NormalizedChunks, + shape: ShapeType | None = None, limit: int | None = None, - dtype: _DType | None = None, - previous_chunks: _NormalizedChunks | None = None, - ) -> _NormalizedChunks: + dtype: DType | None = None, + previous_chunks: NormalizedChunks | None = None, + ) -> NormalizedChunks: """ Normalize given chunking pattern into an explicit tuple of tuples representation. @@ -300,7 +300,7 @@ def normalize_chunks( @abstractmethod def from_array( - self, data: duckarray[Any, Any], chunks: _Chunks, **kwargs: Any + self, data: duckarray[Any, Any], chunks: Chunks, **kwargs: Any ) -> T_ChunkedArray: """ Create a chunked array from a non-chunked numpy-like array. @@ -327,7 +327,7 @@ def from_array( def rechunk( self, data: T_ChunkedArray, - chunks: _NormalizedChunks | tuple[int, ...] | _Chunks, + chunks: NormalizedChunks | tuple[int, ...] | Chunks, **kwargs: Any, ) -> Any: """ @@ -365,7 +365,7 @@ def rechunk( @abstractmethod def compute( self, *data: T_ChunkedArray | Any, **kwargs: Any - ) -> tuple[np.ndarray[Any, _DType_co], ...]: + ) -> tuple[np.ndarray[Any, DType_co], ...]: """ Computes one or more chunked arrays, returning them as eager numpy arrays. @@ -441,7 +441,7 @@ def reduction( combine_func: Callable[..., Any] | None = None, aggregate_func: Callable[..., Any] | None = None, axis: int | Sequence[int] | None = None, - dtype: _DType_co | None = None, + dtype: DType_co | None = None, keepdims: bool = False, ) -> T_ChunkedArray: """ @@ -490,7 +490,7 @@ def scan( ident: float, arr: T_ChunkedArray, axis: int | None = None, - dtype: _DType_co | None = None, + dtype: DType_co | None = None, **kwargs: Any, ) -> T_ChunkedArray: """ @@ -528,7 +528,7 @@ def apply_gufunc( *args: Any, axes: Sequence[tuple[int, ...]] | None = None, keepdims: bool = False, - output_dtypes: Sequence[_DType_co] | None = None, + output_dtypes: Sequence[DType_co] | None = None, vectorize: bool | None = None, **kwargs: Any, ) -> Any: @@ -611,7 +611,7 @@ def map_blocks( self, func: Callable[..., Any], *args: Any, - dtype: _DType_co | None = None, + dtype: DType_co | None = None, chunks: tuple[int, ...] | None = None, drop_axis: int | Sequence[int] | None = None, new_axis: int | Sequence[int] | None = None, @@ -710,7 +710,7 @@ def unify_chunks( self, *args: Any, # can't type this as mypy assumes args are all same type, but dask unify_chunks args alternate types **kwargs: Any, - ) -> tuple[dict[str, _NormalizedChunks], list[T_ChunkedArray]]: + ) -> tuple[dict[str, NormalizedChunks], list[T_ChunkedArray]]: """ Unify chunks across a sequence of arrays. diff --git a/xarray/namedarray/pycompat.py b/xarray/namedarray/pycompat.py index 5832f7cc9e7..2ce66fde48b 100644 --- a/xarray/namedarray/pycompat.py +++ b/xarray/namedarray/pycompat.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: ModType = Literal["dask", "pint", "cupy", "sparse", "cubed", "numbagg"] DuckArrayTypes = tuple[type[Any], ...] # TODO: improve this? maybe Generic - from xarray.namedarray._typing import _DType, _ShapeType, duckarray + from xarray.namedarray._typing import DType, ShapeType, duckarray class DuckArrayModule: @@ -127,7 +127,7 @@ def to_numpy( return data -def to_duck_array(data: Any, **kwargs: dict[str, Any]) -> duckarray[_ShapeType, _DType]: +def to_duck_array(data: Any, **kwargs: dict[str, Any]) -> duckarray[ShapeType, DType]: from xarray.core.indexing import ( ExplicitlyIndexed, ImplicitToExplicitIndexingAdapter, @@ -149,7 +149,7 @@ def to_duck_array(data: Any, **kwargs: dict[str, Any]) -> duckarray[_ShapeType, async def async_to_duck_array( data: Any, **kwargs: dict[str, Any] -) -> duckarray[_ShapeType, _DType]: +) -> duckarray[ShapeType, DType]: from xarray.core.indexing import ( ExplicitlyIndexed, ImplicitToExplicitIndexingAdapter, diff --git a/xarray/namedarray/utils.py b/xarray/namedarray/utils.py index 3490a76aa8d..30f05f8a1f3 100644 --- a/xarray/namedarray/utils.py +++ b/xarray/namedarray/utils.py @@ -4,7 +4,7 @@ import itertools import sys import warnings -from collections.abc import Hashable, Iterable, Iterator, Mapping +from collections.abc import Iterable, Iterator, Mapping from functools import lru_cache from numbers import Number from typing import TYPE_CHECKING, Any, TypeVar, cast @@ -12,9 +12,10 @@ import numpy as np from packaging.version import Version -from xarray.namedarray._typing import ErrorOptionsWithWarn, _DimsLike +from xarray.namedarray._typing import ErrorHandlingWithWarn if TYPE_CHECKING: + from types import EllipsisType from typing import TypeGuard from numpy.typing import NDArray @@ -27,7 +28,7 @@ DaskCollection: Any = NDArray # type: ignore[no-redef] from xarray.core.types import T_ChunkDim - from xarray.namedarray._typing import DuckArray, _Dim, duckarray + from xarray.namedarray._typing import DimType, DuckArray, duckarray from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint @@ -109,10 +110,10 @@ def is_dict_like(value: Any) -> TypeGuard[Mapping[Any, Any]]: def drop_missing_dims( - supplied_dims: Iterable[_Dim], - dims: Iterable[_Dim], - missing_dims: ErrorOptionsWithWarn, -) -> _DimsLike: + supplied_dims: Iterable[DimType | EllipsisType], + dims: Iterable[DimType], + missing_dims: ErrorHandlingWithWarn, +) -> tuple[DimType | EllipsisType, ...]: """Depending on the setting of missing_dims, drop any dimensions from supplied_dims that are not present in dims. @@ -122,39 +123,33 @@ def drop_missing_dims( dims : Iterable of Hashable missing_dims : {"raise", "warn", "ignore"} """ + supplied_dims_tuple = tuple(supplied_dims) + supplied_dims_set = {val for val in supplied_dims_tuple if val is not ...} + dims_set = set(dims) + + if missing_dims in ("raise", "warn"): + if invalid := supplied_dims_set - dims_set: + msg = f"Dimensions {invalid} do not exist. Expected one or more of {dims}" + if missing_dims == "raise": + raise ValueError(msg) + else: + warnings.warn(msg, stacklevel=2) - if missing_dims == "raise": - supplied_dims_set = {val for val in supplied_dims if val is not ...} - if invalid := supplied_dims_set - set(dims): - raise ValueError( - f"Dimensions {invalid} do not exist. Expected one or more of {dims}" - ) - - return supplied_dims - - elif missing_dims == "warn": - if invalid := set(supplied_dims) - set(dims): - warnings.warn( - f"Dimensions {invalid} do not exist. Expected one or more of {dims}", - stacklevel=2, - ) - - return [val for val in supplied_dims if val in dims or val is ...] - - elif missing_dims == "ignore": - return [val for val in supplied_dims if val in dims or val is ...] + return supplied_dims_tuple - else: + elif missing_dims != "ignore": raise ValueError( f"Unrecognised option {missing_dims} for missing_dims argument" ) + return tuple(d for d in supplied_dims_tuple if d in dims_set or d is ...) + def infix_dims( - dims_supplied: Iterable[_Dim], - dims_all: Iterable[_Dim], - missing_dims: ErrorOptionsWithWarn = "raise", -) -> Iterator[_Dim]: + dims_supplied: Iterable[DimType | EllipsisType], + dims_all: Iterable[DimType], + missing_dims: ErrorHandlingWithWarn = "raise", +) -> Iterator[DimType]: """ Resolves a supplied list containing an ellipsis representing other items, to a generator with the 'realized' list of all items @@ -173,7 +168,10 @@ def infix_dims( else: yield d else: - existing_dims = drop_missing_dims(dims_supplied, dims_all, missing_dims) + existing_dims = cast( + tuple[DimType, ...], + drop_missing_dims(dims_supplied, dims_all, missing_dims), + ) if set(existing_dims) ^ set(dims_all): raise ValueError( f"{dims_supplied} must be a permuted list of {dims_all}, unless `...` is included" @@ -182,14 +180,14 @@ def infix_dims( def either_dict_or_kwargs( - pos_kwargs: Mapping[Any, T] | None, + pos_kwargs: Mapping[K, T] | None, kw_kwargs: Mapping[str, T], func_name: str, -) -> Mapping[Hashable, T]: +) -> Mapping[K, T]: if pos_kwargs is None or pos_kwargs == {}: # Need an explicit cast to appease mypy due to invariance; see # https://github.com/python/mypy/issues/6228 - return cast(Mapping[Hashable, T], kw_kwargs) + return cast(Mapping[K, T], kw_kwargs) if not is_dict_like(pos_kwargs): raise ValueError(f"the first argument to .{func_name} must be a dictionary") diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index 09ff1795ef4..3ef782f8f88 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -12,10 +12,10 @@ from xarray.core.indexing import ExplicitlyIndexed from xarray.namedarray._typing import ( + DType_co, + ShapeType_co, _arrayfunction_or_api, _default, - _DType_co, - _ShapeType_co, ) from xarray.namedarray.core import NamedArray, from_array from xarray.namedarray.utils import fake_target_chunksize @@ -27,35 +27,35 @@ from numpy.typing import ArrayLike, DTypeLike, NDArray from xarray.namedarray._typing import ( + AttrsLike, Default, + DimsLike, + DType, DuckArray, - _AttrsLike, + IndexKeyLike, + IntOrUnknown, + Shape, + ShapeLike, _Dim, - _DimsLike, - _DType, - _IndexKeyLike, - _IntOrUnknown, - _Shape, - _ShapeLike, duckarray, ) -class CustomArrayBase(Generic[_ShapeType_co, _DType_co]): - def __init__(self, array: duckarray[Any, _DType_co]) -> None: - self.array: duckarray[Any, _DType_co] = array +class CustomArrayBase(Generic[ShapeType_co, DType_co]): + def __init__(self, array: duckarray[Any, DType_co]) -> None: + self.array: duckarray[Any, DType_co] = array @property - def dtype(self) -> _DType_co: + def dtype(self) -> DType_co: return self.array.dtype @property - def shape(self) -> _Shape: + def shape(self) -> Shape: return self.array.shape class CustomArray( - CustomArrayBase[_ShapeType_co, _DType_co], Generic[_ShapeType_co, _DType_co] + CustomArrayBase[ShapeType_co, DType_co], Generic[ShapeType_co, DType_co] ): def __array__( self, dtype: DTypeLike | None = None, /, *, copy: bool | None = None @@ -67,13 +67,13 @@ def __array__( class CustomArrayIndexable( - CustomArrayBase[_ShapeType_co, _DType_co], + CustomArrayBase[ShapeType_co, DType_co], ExplicitlyIndexed, - Generic[_ShapeType_co, _DType_co], + Generic[ShapeType_co, DType_co], ): def __getitem__( - self, key: _IndexKeyLike | CustomArrayIndexable[Any, Any], / - ) -> CustomArrayIndexable[Any, _DType_co]: + self, key: IndexKeyLike | CustomArrayIndexable[Any, Any], / + ) -> CustomArrayIndexable[Any, DType_co]: if isinstance(key, CustomArrayIndexable): if isinstance(key.array, type(self.array)): # TODO: key.array is duckarray here, can it be narrowed down further? @@ -88,9 +88,9 @@ def __array_namespace__(self) -> ModuleType: return np -def check_duck_array_typevar(a: duckarray[Any, _DType]) -> duckarray[Any, _DType]: +def check_duck_array_typevar(a: duckarray[Any, DType]) -> duckarray[Any, DType]: # Mypy checks a is valid: - b: duckarray[Any, _DType] = a + b: duckarray[Any, DType] = a # Runtime check if valid: if isinstance(b, _arrayfunction_or_api): @@ -220,7 +220,7 @@ def test_init(self, expected: Any) -> None: ) def test_from_array( self, - dims: _DimsLike, + dims: DimsLike, data: ArrayLike, expected: np.ndarray[Any, Any], raise_error: bool, @@ -404,30 +404,30 @@ def test_new_namedarray(self) -> None: assert narr_int.dtype == dtype_int class Variable( - NamedArray[_ShapeType_co, _DType_co], Generic[_ShapeType_co, _DType_co] + NamedArray[ShapeType_co, DType_co], Generic[ShapeType_co, DType_co] ): @overload def _new( self, - dims: _DimsLike | Default = ..., - data: duckarray[Any, _DType] = ..., - attrs: _AttrsLike | Default = ..., - ) -> Variable[Any, _DType]: ... + dims: DimsLike | Default = ..., + data: duckarray[Any, DType] = ..., + attrs: AttrsLike | Default = ..., + ) -> Variable[Any, DType]: ... @overload def _new( self, - dims: _DimsLike | Default = ..., + dims: DimsLike | Default = ..., data: Default = ..., - attrs: _AttrsLike | Default = ..., - ) -> Variable[_ShapeType_co, _DType_co]: ... + attrs: AttrsLike | Default = ..., + ) -> Variable[ShapeType_co, DType_co]: ... def _new( self, - dims: _DimsLike | Default = _default, - data: duckarray[Any, _DType] | Default = _default, - attrs: _AttrsLike | Default = _default, - ) -> Variable[Any, _DType] | Variable[_ShapeType_co, _DType_co]: + dims: DimsLike | Default = _default, + data: duckarray[Any, DType] | Default = _default, + attrs: AttrsLike | Default = _default, + ) -> Variable[Any, DType] | Variable[ShapeType_co, DType_co]: dims_ = copy.copy(self._dims) if dims is _default else dims attrs_: Mapping[Any, Any] | None @@ -438,7 +438,7 @@ def _new( if data is _default: return type(self)(dims_, copy.copy(self._data), attrs_) - cls_ = cast("type[Variable[Any, _DType]]", type(self)) + cls_ = cast("type[Variable[Any, DType]]", type(self)) return cls_(dims_, data, attrs_) var_float: Variable[Any, np.dtype[np.float32]] @@ -465,30 +465,30 @@ def test_replace_namedarray(self) -> None: assert narr_float2.dtype == dtype_float class Variable( - NamedArray[_ShapeType_co, _DType_co], Generic[_ShapeType_co, _DType_co] + NamedArray[ShapeType_co, DType_co], Generic[ShapeType_co, DType_co] ): @overload def _new( self, - dims: _DimsLike | Default = ..., - data: duckarray[Any, _DType] = ..., - attrs: _AttrsLike | Default = ..., - ) -> Variable[Any, _DType]: ... + dims: DimsLike | Default = ..., + data: duckarray[Any, DType] = ..., + attrs: AttrsLike | Default = ..., + ) -> Variable[Any, DType]: ... @overload def _new( self, - dims: _DimsLike | Default = ..., + dims: DimsLike | Default = ..., data: Default = ..., - attrs: _AttrsLike | Default = ..., - ) -> Variable[_ShapeType_co, _DType_co]: ... + attrs: AttrsLike | Default = ..., + ) -> Variable[ShapeType_co, DType_co]: ... def _new( self, - dims: _DimsLike | Default = _default, - data: duckarray[Any, _DType] | Default = _default, - attrs: _AttrsLike | Default = _default, - ) -> Variable[Any, _DType] | Variable[_ShapeType_co, _DType_co]: + dims: DimsLike | Default = _default, + data: duckarray[Any, DType] | Default = _default, + attrs: AttrsLike | Default = _default, + ) -> Variable[Any, DType] | Variable[ShapeType_co, DType_co]: dims_ = copy.copy(self._dims) if dims is _default else dims attrs_: Mapping[Any, Any] | None @@ -499,7 +499,7 @@ def _new( if data is _default: return type(self)(dims_, copy.copy(self._data), attrs_) - cls_ = cast("type[Variable[Any, _DType]]", type(self)) + cls_ = cast("type[Variable[Any, DType]]", type(self)) return cls_(dims_, data, attrs_) var_float: Variable[Any, np.dtype[np.float32]] @@ -523,8 +523,8 @@ def test_expand_dims( target: NamedArray[Any, np.dtype[np.float32]], dim: _Dim | Default, expected_ndim: int, - expected_shape: _ShapeLike, - expected_dims: _DimsLike, + expected_shape: ShapeLike, + expected_dims: DimsLike, ) -> None: result = target.expand_dims(dim=dim) assert result.ndim == expected_ndim @@ -542,8 +542,8 @@ def test_expand_dims( def test_permute_dims( self, target: NamedArray[Any, np.dtype[np.float32]], - dims: _DimsLike, - expected_sizes: dict[_Dim, _IntOrUnknown], + dims: DimsLike, + expected_sizes: dict[_Dim, IntOrUnknown], ) -> None: actual = target.permute_dims(*dims) assert actual.sizes == expected_sizes diff --git a/xarray/tests/test_parallelcompat.py b/xarray/tests/test_parallelcompat.py index 4bb30649676..439885fbf98 100644 --- a/xarray/tests/test_parallelcompat.py +++ b/xarray/tests/test_parallelcompat.py @@ -8,7 +8,7 @@ from xarray import set_options from xarray.core.types import T_Chunks, T_DuckArray, T_NormalizedChunks -from xarray.namedarray._typing import _Chunks +from xarray.namedarray._typing import Chunks from xarray.namedarray.daskmanager import DaskManager from xarray.namedarray.parallelcompat import ( KNOWN_CHUNKMANAGERS, @@ -81,7 +81,7 @@ def normalize_chunks( return normalize_chunks(chunks, shape, limit, dtype, previous_chunks) def from_array( - self, data: T_DuckArray | np.typing.ArrayLike, chunks: _Chunks, **kwargs + self, data: T_DuckArray | np.typing.ArrayLike, chunks: Chunks, **kwargs ) -> DummyChunkedArray: from dask import array as da