Skip to content

Commit fcf3901

Browse files
FIX: Add overloads for DataArray.argmin and DataArray.argmax on dim parameter (pydata#11233)
* Fix mypy overload-overlap error for argmin and argmax - Move `# type: ignore[overload-overlap]` to the @overload decorator - Add whats-new entry * Fix placement type: ignore comments for argmin and argmax * Fix indentation in whats-new.rst --------- Co-authored-by: Michael Niklas <mick.niklas@gmail.com>
1 parent 8ac5ca0 commit fcf3901

2 files changed

Lines changed: 45 additions & 0 deletions

File tree

doc/whats-new.rst

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,10 @@ Internal Changes
135135
- Remove ``setup.py`` file (:pull:`11261`).
136136
By `Nick Hodgskin <https://github.com/VeckoTheGecko>`_.
137137

138+
- Add :func:`typing.overload` decorators to :py:meth:`DataArray.argmin` and :py:meth:`DataArray.argmax`
139+
to narrow return type based on ``dim`` parameter (:issue:`10893` :pull:`11233`).
140+
By `Amartya Anand <https://github.com/SurfyPenguin>`_.
141+
138142
.. _whats-new.2026.02.0:
139143

140144
v2026.02.0 (Feb 13, 2026)

xarray/core/dataarray.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import warnings
66
from collections.abc import (
77
Callable,
8+
Collection,
89
Hashable,
910
Iterable,
1011
Mapping,
@@ -6187,6 +6188,26 @@ def idxmax(
61876188
keep_attrs=keep_attrs,
61886189
)
61896190

6191+
@overload
6192+
def argmin( # type: ignore[overload-overlap]
6193+
self,
6194+
dim: str,
6195+
*,
6196+
axis: int | None = None,
6197+
keep_attrs: bool | None = None,
6198+
skipna: bool | None = None,
6199+
) -> Self: ...
6200+
6201+
@overload
6202+
def argmin(
6203+
self,
6204+
dim: Collection[Hashable] | EllipsisType | None = None,
6205+
*,
6206+
axis: int | None = None,
6207+
keep_attrs: bool | None = None,
6208+
skipna: bool | None = None,
6209+
) -> dict[Hashable, Self]: ...
6210+
61906211
def argmin(
61916212
self,
61926213
dim: Dims = None,
@@ -6288,6 +6309,26 @@ def argmin(
62886309
else:
62896310
return self._replace_maybe_drop_dims(result)
62906311

6312+
@overload
6313+
def argmax( # type: ignore[overload-overlap]
6314+
self,
6315+
dim: str,
6316+
*,
6317+
axis: int | None = None,
6318+
keep_attrs: bool | None = None,
6319+
skipna: bool | None = None,
6320+
) -> Self: ...
6321+
6322+
@overload
6323+
def argmax(
6324+
self,
6325+
dim: Collection[Hashable] | EllipsisType | None = None,
6326+
*,
6327+
axis: int | None = None,
6328+
keep_attrs: bool | None = None,
6329+
skipna: bool | None = None,
6330+
) -> dict[Hashable, Self]: ...
6331+
62916332
def argmax(
62926333
self,
62936334
dim: Dims = None,

0 commit comments

Comments
 (0)