Skip to content

Add overloads for DataArray.argmin and DataArray.argmax on dim parameter#11233

Open
SurfyPenguin wants to merge 4 commits intopydata:mainfrom
SurfyPenguin:fix/argmin-argmax-overloads
Open

Add overloads for DataArray.argmin and DataArray.argmax on dim parameter#11233
SurfyPenguin wants to merge 4 commits intopydata:mainfrom
SurfyPenguin:fix/argmin-argmax-overloads

Conversation

@SurfyPenguin
Copy link

What this does

Adds @overload decorators to DataArray.argmin and DataArray.argmax so type checkers can narrow the return type based on the dim argument.

  • dim: str -> Self
  • dim: Collection[Hashable] | EllipsisType | None = None -> dict[Hashable, Self]

Additional info

While testing, I noticed this FutureWarning:

FutureWarning: Behaviour of argmin/argmax with neither dim nor axis argument will change to return a dict of indices of each dimension. To get a single, flat index, please use np.argmin(da.data) or np.argmax(da.data) instead of da.argmin() or da.argmax().

The overload reflects the intended behavior i.e. dim=None is typed to return dict[Hashable, Self] rather than Self.

@welcome
Copy link

welcome bot commented Mar 15, 2026

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

@SurfyPenguin SurfyPenguin force-pushed the fix/argmin-argmax-overloads branch 3 times, most recently from 8c91f37 to 471352b Compare March 15, 2026 21:56
- Move `# type: ignore[overload-overlap]` to the @overload decorator
- Add whats-new entry
@SurfyPenguin SurfyPenguin force-pushed the fix/argmin-argmax-overloads branch from 471352b to f950225 Compare March 15, 2026 21:59
@SurfyPenguin
Copy link
Author

Closing this PR for now.
After syncing with main and re-running mypy. I realized my overload pattern for DataArray.argmin and DataArray.argmax reproduces the exact overlap issue tracked in #10893 .

@SurfyPenguin SurfyPenguin deleted the fix/argmin-argmax-overloads branch March 15, 2026 22:38
Copy link
Collaborator

@headtr1ck headtr1ck left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This approach is actually correct. The ignore is important and it will work exactly as intended.

@SurfyPenguin SurfyPenguin restored the fix/argmin-argmax-overloads branch March 16, 2026 20:07
@SurfyPenguin SurfyPenguin reopened this Mar 16, 2026
@SurfyPenguin
Copy link
Author

@headtr1ck , thanks for the clarification. I was not sure if type: ignore[overload-overlap] was the right thing to do here. Good to know its intentional. I'm re-opening the PR.

@headtr1ck
Copy link
Collaborator

Yes it's correct. Mypy correctly complains because technically it is overlapping because str is an IT table of hashable. But if you ignore this, it will correctly infer that str will return Self while e.g. a list returns something else (it will start looking through the overloads by their order and pick the first match).

@headtr1ck headtr1ck added plan to merge Final call for comments topic-typing labels Mar 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

plan to merge Final call for comments topic-typing

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Overload DataArray's argmin and argmax methods on dim parameter

2 participants