Skip to content

TYP: Fix np.interp signature for scalar types #27869

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jan 9, 2025
Merged

TYP: Fix np.interp signature for scalar types #27869

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 82 additions & 15 deletions numpy/lib/_function_base_impl.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from collections.abc import Sequence, Iterator, Callable, Iterable
from collections.abc import Sequence, Callable, Iterable
from typing import (
Concatenate,
Literal as L,
Expand All @@ -15,8 +15,9 @@ from typing import (
)
from typing_extensions import deprecated

import numpy as np
from numpy import (
vectorize as vectorize,
Copy link
Member

Choose a reason for hiding this comment

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

Remind me of why vectorize as vectorize makes sense. I've forgotten.

Copy link
Member

Choose a reason for hiding this comment

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

In this case it doesn't, but if there's no __all__, an import spatula from spatula will implicitly export spatula as a public member of that module. It's almost always better to use __all__, but mypy's stubgen tool used to (since last month) generate stubs with such implicit exports.

vectorize,
generic,
integer,
floating,
Expand All @@ -35,19 +36,22 @@ from numpy._typing import (
NDArray,
ArrayLike,
DTypeLike,
_ShapeLike,
_ScalarLike_co,
_DTypeLike,
_ArrayLike,
_DTypeLike,
_ShapeLike,
_ArrayLikeBool_co,
_ArrayLikeInt_co,
_ArrayLikeFloat_co,
_ArrayLikeComplex_co,
_ArrayLikeNumber_co,
_ArrayLikeTD64_co,
_ArrayLikeDT64_co,
_ArrayLikeObject_co,
_FloatLike_co,
_ComplexLike_co,
_NumberLike_co,
_ScalarLike_co,
_NestedSequence
)

__all__ = [
Expand Down Expand Up @@ -303,24 +307,87 @@ def diff(
append: ArrayLike = ...,
) -> NDArray[Any]: ...

@overload
@overload # float scalar
def interp(
x: _ArrayLikeFloat_co,
x: _FloatLike_co,
xp: _ArrayLikeFloat_co,
fp: _ArrayLikeFloat_co,
left: _FloatLike_co | None = None,
right: _FloatLike_co | None = None,
period: _FloatLike_co | None = None,
) -> float64: ...
@overload # float array
def interp(
x: NDArray[floating | integer | np.bool] | _NestedSequence[_FloatLike_co],
xp: _ArrayLikeFloat_co,
fp: _ArrayLikeFloat_co,
left: None | _FloatLike_co = ...,
right: None | _FloatLike_co = ...,
period: None | _FloatLike_co = ...,
left: _FloatLike_co | None = None,
Copy link
Member

Choose a reason for hiding this comment

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

Curious about this usage of None = None.

Copy link
Member

@jorenham jorenham Jan 9, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

Two things changed here:

  1. The None should always be the last type within a union (according to the typing spec style guide, ruff, and flak8-pyi), so I tend to correct that in the places where I change things in the stubs (not the best for the diffs, but I can't really help myself 😅).
  2. You noted before that annotations are a form of documentation. And when documenting functions, you also document the defaults. Using = ... in stubs was done by mypy's stubgen (for unknown reasons), and was therefore assumed to be the way to go. But for documentation, and also for things like IDE suggestions and type-checker error messages, it's way more helpful if the actual default value is used, instead of a .... I believe the typing spec recently also started recommending doing this.

right: _FloatLike_co | None = None,
period: _FloatLike_co | None = None,
) -> NDArray[float64]: ...
@overload
@overload # float scalar or array
def interp(
x: _ArrayLikeFloat_co,
xp: _ArrayLikeFloat_co,
fp: _ArrayLikeComplex_co,
left: None | _ComplexLike_co = ...,
right: None | _ComplexLike_co = ...,
period: None | _FloatLike_co = ...,
fp: _ArrayLikeFloat_co,
left: _FloatLike_co | None = None,
right: _FloatLike_co | None = None,
period: _FloatLike_co | None = None,
) -> NDArray[float64] | float64: ...
@overload # complex scalar
def interp(
x: _FloatLike_co,
xp: _ArrayLikeFloat_co,
fp: _ArrayLike[complexfloating],
left: _NumberLike_co | None = None,
right: _NumberLike_co | None = None,
period: _FloatLike_co | None = None,
) -> complex128: ...
@overload # complex or float scalar
def interp(
x: _FloatLike_co,
xp: _ArrayLikeFloat_co,
fp: Sequence[complex | complexfloating],
left: _NumberLike_co | None = None,
right: _NumberLike_co | None = None,
period: _FloatLike_co | None = None,
) -> complex128 | float64: ...
@overload # complex array
def interp(
x: NDArray[floating | integer | np.bool] | _NestedSequence[_FloatLike_co],
xp: _ArrayLikeFloat_co,
fp: _ArrayLike[complexfloating],
left: _NumberLike_co | None = None,
right: _NumberLike_co | None = None,
period: _FloatLike_co | None = None,
) -> NDArray[complex128]: ...
@overload # complex or float array
def interp(
x: NDArray[floating | integer | np.bool] | _NestedSequence[_FloatLike_co],
xp: _ArrayLikeFloat_co,
fp: Sequence[complex | complexfloating],
left: _NumberLike_co | None = None,
right: _NumberLike_co | None = None,
period: _FloatLike_co | None = None,
) -> NDArray[complex128 | float64]: ...
@overload # complex scalar or array
def interp(
x: _ArrayLikeFloat_co,
xp: _ArrayLikeFloat_co,
fp: _ArrayLike[complexfloating],
left: _NumberLike_co | None = None,
right: _NumberLike_co | None = None,
period: _FloatLike_co | None = None,
) -> NDArray[complex128] | complex128: ...
@overload # complex or float scalar or array
def interp(
x: _ArrayLikeFloat_co,
xp: _ArrayLikeFloat_co,
fp: _ArrayLikeNumber_co,
left: _NumberLike_co | None = None,
right: _NumberLike_co | None = None,
period: _FloatLike_co | None = None,
) -> NDArray[complex128 | float64] | complex128 | float64: ...

@overload
def angle(z: _ComplexLike_co, deg: bool = ...) -> floating[Any]: ...
Expand Down
9 changes: 9 additions & 0 deletions numpy/typing/tests/data/reveal/lib_function_base.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,15 @@ assert_type(np.diff("bob", n=0), str)
assert_type(np.diff(AR_f8, axis=0), npt.NDArray[Any])
assert_type(np.diff(AR_LIKE_f8, prepend=1.5), npt.NDArray[Any])

assert_type(np.interp(1, [1], AR_f8), np.float64)
assert_type(np.interp(1, [1], [1]), np.float64)
assert_type(np.interp(1, [1], AR_c16), np.complex128)
assert_type(np.interp(1, [1], [1j]), np.complex128) # pyright correctly infers `complex128 | float64`
assert_type(np.interp([1], [1], AR_f8), npt.NDArray[np.float64])
assert_type(np.interp([1], [1], [1]), npt.NDArray[np.float64])
assert_type(np.interp([1], [1], AR_c16), npt.NDArray[np.complex128])
assert_type(np.interp([1], [1], [1j]), npt.NDArray[np.complex128]) # pyright correctly infers `NDArray[complex128 | float64]`

assert_type(np.angle(f8), np.floating[Any])
assert_type(np.angle(AR_f8), npt.NDArray[np.floating[Any]])
assert_type(np.angle(AR_c16, deg=True), npt.NDArray[np.floating[Any]])
Expand Down
Loading