Skip to content

Commit 7131c9b

Browse files
jorenhamjorenham
authored
✅ fix failing type-tests on numpy==2.1.* (#788)
2 parents b554329 + 53c2a60 commit 7131c9b

File tree

5 files changed

+58
-43
lines changed

5 files changed

+58
-43
lines changed

‎scipy-stubs/integrate/_ivp/ivp.pyi‎

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,15 @@ _ScalarT = TypeVar("_ScalarT", bound=npc.number | np.bool_)
1717
_Inexact64T = TypeVar("_Inexact64T", bound=np.float64 | np.complex128)
1818
_Inexact64T_co = TypeVar("_Inexact64T_co", bound=np.float64 | np.complex128, default=np.float64 | np.complex128, covariant=True)
1919

20-
_FuncSol: TypeAlias = Callable[[np.float64], onp.ArrayND[_Inexact64T]]
21-
_FuncEvent: TypeAlias = Callable[[np.float64, onp.ArrayND[_Inexact64T], *_Ts], float]
20+
# numpy <2.2 workarounds
21+
_Float: TypeAlias = np.float64 | float
22+
_FloatT = TypeVar("_FloatT", bound=_Float)
23+
24+
_FuncSol: TypeAlias = Callable[[np.float64], onp.ArrayND[_Inexact64T]] | Callable[[float], onp.ArrayND[_Inexact64T]]
25+
_FuncEvent: TypeAlias = (
26+
Callable[[np.float64, onp.ArrayND[_Inexact64T], *_Ts], _Float]
27+
| Callable[[float, onp.ArrayND[_Inexact64T], *_Ts], _Float]
28+
) # fmt: skip
2229
_Events: TypeAlias = Sequence[_FuncEvent[_Inexact64T, *_Ts]] | _FuncEvent[_Inexact64T, *_Ts]
2330

2431
_Int1D: TypeAlias = onp.Array1D[np.int_]
@@ -27,6 +34,9 @@ _Float2D: TypeAlias = onp.Array2D[np.float64]
2734
_Complex1D: TypeAlias = onp.Array1D[np.complex128]
2835
_Complex2D: TypeAlias = onp.Array2D[np.complex128]
2936

37+
_ToFloatMax1D: TypeAlias = onp.ToFloat1D | onp.ToFloat
38+
_ToComplexMax1D: TypeAlias = onp.ToComplex1D | onp.ToComplex
39+
3040
_Sparse2D: TypeAlias = _spbase[_ScalarT, tuple[int, int]] | sparray[_ScalarT, tuple[int, int]] | spmatrix[_ScalarT]
3141
_ToJac: TypeAlias = onp.ToArray2D[complex, npc.inexact] | _Sparse2D[npc.inexact]
3242

@@ -75,10 +85,10 @@ def handle_events(
7585
) -> tuple[_Int1D, _Float1D, bool]: ...
7686
def find_active_events(g: onp.ToFloat1D, g_new: onp.ToFloat1D, direction: onp.ArrayND[np.float64]) -> _Int1D: ...
7787

78-
#
88+
# NOTE: The *free* `_FloatT` type variable works around `float64` not being a subtype of `float` on `numpy <2.2`.
7989
@overload # float, vectorized=False (default), args=None (default)
8090
def solve_ivp(
81-
fun: Callable[[np.float64, _Float1D], onp.ToFloat1D | float],
91+
fun: Callable[[_FloatT, _Float1D], _ToFloatMax1D],
8292
t_span: Sequence[float],
8393
y0: onp.ToFloat1D,
8494
method: _IVPMethod = "RK45",
@@ -91,7 +101,7 @@ def solve_ivp(
91101
) -> OdeResult[np.float64]: ...
92102
@overload # float, vectorized=False (default), args=<given>
93103
def solve_ivp(
94-
fun: Callable[[np.float64, _Float1D, *_Ts], onp.ToFloat1D | float],
104+
fun: Callable[[_FloatT, _Float1D, *_Ts], _ToFloatMax1D],
95105
t_span: Sequence[float],
96106
y0: onp.ToFloat1D,
97107
method: _IVPMethod = "RK45",
@@ -133,7 +143,7 @@ def solve_ivp(
133143
) -> OdeResult[np.float64]: ...
134144
@overload # complex, vectorized=False (default), args=None (default)
135145
def solve_ivp(
136-
fun: Callable[[np.float64, _Complex1D], onp.ToComplex1D | complex],
146+
fun: Callable[[_FloatT, _Complex1D], _ToComplexMax1D],
137147
t_span: Sequence[float],
138148
y0: onp.ToComplex1D,
139149
method: _IVPMethod = "RK45",
@@ -146,7 +156,7 @@ def solve_ivp(
146156
) -> OdeResult[np.complex128]: ...
147157
@overload # complex, vectorized=False (default), args=<given>
148158
def solve_ivp(
149-
fun: Callable[[np.float64, _Complex1D, *_Ts], onp.ToComplex1D | complex],
159+
fun: Callable[[_FloatT, _Complex1D, *_Ts], _ToComplexMax1D],
150160
t_span: Sequence[float],
151161
y0: onp.ToComplex1D,
152162
method: _IVPMethod = "RK45",

‎scipy-stubs/interpolate/_polyint.pyi‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -118,22 +118,6 @@ class KroghInterpolator(_Interpolator1DWithDerivatives[_YT_co], Generic[_YT_co,
118118
n: Final[int] # undocumented
119119
r: Final[int] # undocumented
120120

121-
@overload # xi: i64, yi: f64
122-
def __init__(
123-
self: KroghInterpolator[np.float64, np.float64], /, xi: onp.ToJustInt64_1D, yi: onp.ToFloatND, axis: int = 0
124-
) -> None: ...
125-
@overload # xi: i64, yi: c128
126-
def __init__(
127-
self: KroghInterpolator[np.complex128, np.float64], /, xi: onp.ToJustInt64_1D, yi: onp.ToJustComplexND, axis: int = 0
128-
) -> None: ...
129-
@overload # xi: f64, yi: f64
130-
def __init__(
131-
self: KroghInterpolator[np.float64, np.float64], /, xi: onp.ToJustFloat64_1D, yi: onp.ToFloatND, axis: int = 0
132-
) -> None: ...
133-
@overload # xi: f64, yi: c128
134-
def __init__(
135-
self: KroghInterpolator[np.complex128, np.float64], /, xi: onp.ToJustFloat64_1D, yi: onp.ToJustComplexND, axis: int = 0
136-
) -> None: ...
137121
@overload # xi: T, yi: f64
138122
def __init__(
139123
self: KroghInterpolator[np.float64, _XT],
@@ -150,6 +134,22 @@ class KroghInterpolator(_Interpolator1DWithDerivatives[_YT_co], Generic[_YT_co,
150134
yi: onp.ToJustComplexND,
151135
axis: int = 0,
152136
) -> None: ...
137+
@overload # xi: i64, yi: f64
138+
def __init__(
139+
self: KroghInterpolator[np.float64, np.float64], /, xi: onp.ToJustInt64_1D, yi: onp.ToFloatND, axis: int = 0
140+
) -> None: ...
141+
@overload # xi: i64, yi: c128
142+
def __init__(
143+
self: KroghInterpolator[np.complex128, np.float64], /, xi: onp.ToJustInt64_1D, yi: onp.ToJustComplexND, axis: int = 0
144+
) -> None: ...
145+
@overload # xi: f64, yi: f64
146+
def __init__(
147+
self: KroghInterpolator[np.float64, np.float64], /, xi: onp.ToJustFloat64_1D, yi: onp.ToFloatND, axis: int = 0
148+
) -> None: ...
149+
@overload # xi: f64, yi: c128
150+
def __init__(
151+
self: KroghInterpolator[np.complex128, np.float64], /, xi: onp.ToJustFloat64_1D, yi: onp.ToJustComplexND, axis: int = 0
152+
) -> None: ...
153153

154154
#
155155
@override

‎scipy-stubs/stats/_distribution_infrastructure.pyi‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ class _DuckDistributionBase(Protocol[_Tss]):
7474
def __make_distribution_version__(self, /) -> str: ...
7575
@property
7676
def support(self, /) -> _ParameterSpec: ...
77-
def pdf(self, x: float, /, *do_not_use_these: _Tss.args, **parameters: _Tss.kwargs) -> float: ...
77+
def pdf(self, x: float, /, *do_not_use_these: _Tss.args, **parameters: _Tss.kwargs) -> float | np.float64: ...
7878

7979
@final
8080
@type_check_only

‎tests/constants/test_constants.pyi‎

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,26 +5,31 @@ import optype.numpy as onp
55

66
from scipy import constants
77

8+
u64_nd: onp.ArrayND[np.uint64]
9+
f32_nd: onp.ArrayND[np.float32]
10+
f64_nd: onp.ArrayND[np.float64]
11+
c64_nd: onp.ArrayND[np.complex64]
12+
813
assert_type(constants.convert_temperature(1, "Celsius", "Kelvin"), np.float64)
914
assert_type(constants.convert_temperature([1], "Celsius", "Kelvin"), onp.ArrayND[np.float64])
10-
assert_type(constants.convert_temperature(np.array([1], dtype=np.uint64), "Celsius", "Kelvin"), onp.ArrayND[np.float64])
11-
assert_type(constants.convert_temperature(np.array([1], dtype=np.float32), "Celsius", "Kelvin"), onp.ArrayND[np.float32])
12-
assert_type(constants.convert_temperature(np.array([1], dtype=np.float64), "Celsius", "Kelvin"), onp.ArrayND[np.float64])
13-
assert_type(constants.convert_temperature(np.array([1], dtype=np.complex64), "Celsius", "Kelvin"), onp.ArrayND[np.complex64])
15+
assert_type(constants.convert_temperature(u64_nd, "Celsius", "Kelvin"), onp.ArrayND[np.float64])
16+
assert_type(constants.convert_temperature(f32_nd, "Celsius", "Kelvin"), onp.ArrayND[np.float32])
17+
assert_type(constants.convert_temperature(f64_nd, "Celsius", "Kelvin"), onp.ArrayND[np.float64])
18+
assert_type(constants.convert_temperature(c64_nd, "Celsius", "Kelvin"), onp.ArrayND[np.complex64])
1419

1520
assert_type(constants.lambda2nu(1), np.float64)
1621
assert_type(constants.lambda2nu([1]), onp.ArrayND[np.float64])
17-
assert_type(constants.lambda2nu(np.array([1], dtype=np.uint64)), onp.ArrayND[np.float64])
18-
assert_type(constants.lambda2nu(np.array([1], dtype=np.float32)), onp.ArrayND[np.float32])
19-
assert_type(constants.lambda2nu(np.array([1], dtype=np.float64)), onp.ArrayND[np.float64])
20-
assert_type(constants.lambda2nu(np.array([1], dtype=np.complex64)), onp.ArrayND[np.complex64])
22+
assert_type(constants.lambda2nu(u64_nd), onp.ArrayND[np.float64])
23+
assert_type(constants.lambda2nu(f32_nd), onp.ArrayND[np.float32])
24+
assert_type(constants.lambda2nu(f64_nd), onp.ArrayND[np.float64])
25+
assert_type(constants.lambda2nu(c64_nd), onp.ArrayND[np.complex64])
2126

2227
assert_type(constants.nu2lambda(1), np.float64)
2328
assert_type(constants.nu2lambda([1]), onp.ArrayND[np.float64])
24-
assert_type(constants.nu2lambda(np.array([1], dtype=np.uint64)), onp.ArrayND[np.float64])
25-
assert_type(constants.nu2lambda(np.array([1], dtype=np.float32)), onp.ArrayND[np.float32])
26-
assert_type(constants.nu2lambda(np.array([1], dtype=np.float64)), onp.ArrayND[np.float64])
27-
assert_type(constants.nu2lambda(np.array([1], dtype=np.complex64)), onp.ArrayND[np.complex64])
29+
assert_type(constants.nu2lambda(u64_nd), onp.ArrayND[np.float64])
30+
assert_type(constants.nu2lambda(f32_nd), onp.ArrayND[np.float32])
31+
assert_type(constants.nu2lambda(f64_nd), onp.ArrayND[np.float64])
32+
assert_type(constants.nu2lambda(c64_nd), onp.ArrayND[np.complex64])
2833

2934
###
3035

‎tests/integrate/test_solve_ivp.pyi‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,17 @@
1-
from typing import TypeAlias, assert_type, type_check_only
1+
from typing import Any, TypeAlias, assert_type, type_check_only
22

33
import numpy as np
44
import numpy.typing as npt
5+
import optype.numpy as onp
56

67
from scipy.integrate import solve_ivp
78

89
_VecF64: TypeAlias = np.ndarray[tuple[int], np.dtype[np.float64]]
910
_MatF64: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.float64]]
1011
_ArrF64: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.float64]]
11-
_VecC128: TypeAlias = np.ndarray[tuple[int], np.dtype[np.complex128]]
12+
_VecC128: TypeAlias = np.ndarray[tuple[Any], np.dtype[np.complex128]]
1213
_MatC128: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.complex128]]
13-
_ArrC128: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.complex128]]
14+
_ArrC128: TypeAlias = np.ndarray[tuple[Any, ...], np.dtype[np.complex128]]
1415

1516
list_float: list[float] = ...
1617
list_complex: list[complex] = ...
@@ -38,24 +39,23 @@ def upward_cannon(t: np.float64, y: _VecF64) -> list[float]: ...
3839
@type_check_only
3940
def hit_ground(t: np.float64, y: _VecF64) -> np.float64: ...
4041

42+
assert_type(solve_ivp(upward_cannon, list_float, list_float).y, _MatF64)
4143
assert_type(solve_ivp(upward_cannon, list_float, list_float, events=hit_ground).y, _MatF64)
4244
assert_type(solve_ivp(upward_cannon, list_float, list_float, events=hit_ground, args=()).y, _MatF64)
4345
assert_type(solve_ivp(upward_cannon, list_float, list_float, events=hit_ground, dense_output=True).y, _MatF64)
4446

4547
###
4648

4749
@type_check_only
48-
def lotkavolterra(
49-
t: float, z: np.ndarray[tuple[int, ...], np.dtype[np.float64]], a: float, b: float, c: float, d: float
50-
) -> _VecF64: ...
50+
def lotkavolterra(t: float, z: npt.NDArray[np.float64], a: float, b: float, c: float, d: float) -> _VecF64: ...
5151

5252
assert_type(solve_ivp(lotkavolterra, list_float, list_float, args=(1.5, 1, 3, 1)).y, _MatF64)
5353
assert_type(solve_ivp(lotkavolterra, list_float, list_float, args=(1.5, 1, 3, 1), dense_output=True).y, _MatF64)
5454

5555
###
5656

5757
@type_check_only
58-
def deriv_vec(t: float, y: npt.NDArray[np.float64 | np.complex128]) -> npt.NDArray[np.float64 | np.complex128]: ...
58+
def deriv_vec(t: float, y: onp.ArrayND[np.float64 | np.complex128]) -> onp.ArrayND[np.float64 | np.complex128]: ...
5959

6060
assert_type(solve_ivp(deriv_vec, list_float, list_complex).y, _MatC128)
6161
assert_type(solve_ivp(deriv_vec, list_float, vec_c128).y, _MatC128)

0 commit comments

Comments
 (0)