diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 64fddac3460..41679520930 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -135,6 +135,10 @@ Internal Changes - Remove ``setup.py`` file (:pull:`11261`). By `Nick Hodgskin `_. +- Add :func:`typing.overload` decorators to :py:meth:`DataArray.argmin` and :py:meth:`DataArray.argmax` + to narrow return type based on ``dim`` parameter (:issue:`10893` :pull:`11233`). + By `Amartya Anand `_. + .. _whats-new.2026.02.0: v2026.02.0 (Feb 13, 2026) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 0e7ebb5f1bc..c49a41de108 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -5,6 +5,7 @@ import warnings from collections.abc import ( Callable, + Collection, Hashable, Iterable, Mapping, @@ -6187,6 +6188,26 @@ def idxmax( keep_attrs=keep_attrs, ) + @overload + def argmin( # type: ignore[overload-overlap] + self, + dim: str, + *, + axis: int | None = None, + keep_attrs: bool | None = None, + skipna: bool | None = None, + ) -> Self: ... + + @overload + def argmin( + self, + dim: Collection[Hashable] | EllipsisType | None = None, + *, + axis: int | None = None, + keep_attrs: bool | None = None, + skipna: bool | None = None, + ) -> dict[Hashable, Self]: ... + def argmin( self, dim: Dims = None, @@ -6288,6 +6309,26 @@ def argmin( else: return self._replace_maybe_drop_dims(result) + @overload + def argmax( # type: ignore[overload-overlap] + self, + dim: str, + *, + axis: int | None = None, + keep_attrs: bool | None = None, + skipna: bool | None = None, + ) -> Self: ... + + @overload + def argmax( + self, + dim: Collection[Hashable] | EllipsisType | None = None, + *, + axis: int | None = None, + keep_attrs: bool | None = None, + skipna: bool | None = None, + ) -> dict[Hashable, Self]: ... + def argmax( self, dim: Dims = None,