add _DimType_co generic typevar to NamedArray

Also makes private typing names 'public' by removing leading _

Author: headtr1ck

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.
 

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[

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:

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,