Skip to content

Commit 96d3db5

Browse files
jorenhamjorenham
authored
🏷️ fix stubtest errors in numpy.lib._nanfunctions_impl (#244)
1 parent de9e9de commit 96d3db5

File tree

3 files changed

+261
-4
lines changed

3 files changed

+261
-4
lines changed

‎.mypyignore-todo

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ numpy(\..+)?\.floating.as_integer_ratio
88
numpy(\..+)?\.complexfloating.__hash__
99
numpy(\..+)?\.complexfloating.__complex__
1010

11-
numpy(\.lib\._nanfunctions_impl|\.matlib)?\.nan(median|percentile)
1211
numpy(\.lib\._twodim_base_impl)?\.tri(l|u)
1312

1413
numpy\._?core(\._multiarray_umath|\.multiarray)\.error

‎src/numpy-stubs/lib/_function_base_impl.pyi

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -700,7 +700,7 @@ def sinc(x: CoFloating_1nd) -> Array[np.floating]: ...
700700
@overload
701701
def sinc(x: ToComplex_1nd) -> Array[np.complexfloating]: ...
702702

703-
#
703+
# keep in sync with `lib._nanfunctions_impl.nanmedian`
704704
@overload
705705
def median(
706706
a: CoFloating_nd,
@@ -759,6 +759,7 @@ def median(
759759
keepdims: bool = False,
760760
) -> _ArrayT: ...
761761

762+
# keep in sync with `lib._nanfunctions_impl.nanpercentile`
762763
# TODO(jorenham): deprecate interpolation
763764
# TODO(jorenham): deprecate only allow weights if method="inverted_cdf"
764765
@overload

‎src/numpy-stubs/lib/_nanfunctions_impl.pyi

Lines changed: 259 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,20 @@
1-
# NOTE: In reality these functions are not aliases but distinct functions with identical signatures.
1+
from typing import Any, Literal as L, overload
2+
from typing_extensions import TypeVar
3+
4+
import numpy as np
5+
from _numtype import (
6+
Array,
7+
CoComplex_nd,
8+
CoDateTime_nd,
9+
CoFloating_0d,
10+
CoFloating_1nd,
11+
CoFloating_nd,
12+
CoTimeDelta_nd,
13+
ToComplex_nd,
14+
ToDateTime_nd,
15+
ToObject_nd,
16+
ToTimeDelta_nd,
17+
)
218
from numpy._core.fromnumeric import (
319
amax as nanmax,
420
amin as nanmin,
@@ -12,7 +28,10 @@ from numpy._core.fromnumeric import (
1228
sum as nansum,
1329
var as nanvar,
1430
)
15-
from numpy.lib._function_base_impl import median as nanmedian, percentile as nanpercentile, quantile as nanquantile
31+
from numpy._globals import _NoValueType
32+
from numpy._typing import _ShapeLike
33+
34+
from ._function_base_impl import _PercentileMethod
1635

1736
__all__ = [
1837
"nanargmax",
@@ -30,3 +49,241 @@ __all__ = [
3049
"nansum",
3150
"nanvar",
3251
]
52+
53+
###
54+
55+
_ArrayT = TypeVar("_ArrayT", bound=Array)
56+
57+
###
58+
59+
# keep in sync with `lib._function_base_impl.median`
60+
@overload
61+
def nanmedian(
62+
a: CoFloating_nd,
63+
axis: None = None,
64+
out: None = None,
65+
overwrite_input: bool = False,
66+
keepdims: _NoValueType | L[False] = ...,
67+
) -> np.floating: ...
68+
@overload
69+
def nanmedian(
70+
a: ToComplex_nd,
71+
axis: None = None,
72+
out: None = None,
73+
overwrite_input: bool = False,
74+
keepdims: _NoValueType | L[False] = ...,
75+
) -> np.complexfloating: ...
76+
@overload
77+
def nanmedian(
78+
a: ToTimeDelta_nd,
79+
axis: None = None,
80+
out: None = None,
81+
overwrite_input: bool = False,
82+
keepdims: _NoValueType | L[False] = ...,
83+
) -> np.timedelta64: ...
84+
@overload
85+
def nanmedian(
86+
a: ToObject_nd,
87+
axis: None = None,
88+
out: None = None,
89+
overwrite_input: bool = False,
90+
keepdims: _NoValueType | L[False] = ...,
91+
) -> Any: ...
92+
@overload
93+
def nanmedian(
94+
a: CoComplex_nd | CoTimeDelta_nd | ToObject_nd,
95+
axis: _ShapeLike | None = None,
96+
out: None = None,
97+
overwrite_input: bool = False,
98+
keepdims: _NoValueType | bool = ...,
99+
) -> Any: ...
100+
@overload
101+
def nanmedian(
102+
a: CoComplex_nd | CoTimeDelta_nd | ToObject_nd,
103+
axis: _ShapeLike | None,
104+
out: _ArrayT,
105+
overwrite_input: bool = False,
106+
keepdims: _NoValueType | bool = ...,
107+
) -> _ArrayT: ...
108+
@overload
109+
def nanmedian(
110+
a: CoComplex_nd | CoTimeDelta_nd | ToObject_nd,
111+
axis: _ShapeLike | None = None,
112+
*,
113+
out: _ArrayT,
114+
overwrite_input: bool = False,
115+
keepdims: _NoValueType | bool = ...,
116+
) -> _ArrayT: ...
117+
118+
# keep in sync with `lib._function_base_impl.percentile`
119+
@overload
120+
def nanpercentile(
121+
a: CoFloating_nd,
122+
q: CoFloating_0d,
123+
axis: None = None,
124+
out: None = None,
125+
overwrite_input: bool = False,
126+
method: _PercentileMethod = "linear",
127+
keepdims: _NoValueType | L[False] = ...,
128+
*,
129+
weights: CoFloating_1nd | None = None,
130+
interpolation: None = None,
131+
) -> np.floating: ...
132+
@overload
133+
def nanpercentile(
134+
a: CoFloating_nd,
135+
q: CoFloating_1nd,
136+
axis: None = None,
137+
out: None = None,
138+
overwrite_input: bool = False,
139+
method: _PercentileMethod = "linear",
140+
keepdims: _NoValueType | L[False] = ...,
141+
*,
142+
weights: CoFloating_1nd | None = None,
143+
interpolation: None = None,
144+
) -> Array[np.floating]: ...
145+
@overload
146+
def nanpercentile(
147+
a: ToComplex_nd,
148+
q: CoFloating_0d,
149+
axis: None = None,
150+
out: None = None,
151+
overwrite_input: bool = False,
152+
method: _PercentileMethod = "linear",
153+
keepdims: _NoValueType | L[False] = ...,
154+
*,
155+
weights: CoFloating_1nd | None = None,
156+
interpolation: None = None,
157+
) -> np.complexfloating: ...
158+
@overload
159+
def nanpercentile(
160+
a: ToComplex_nd,
161+
q: CoFloating_1nd,
162+
axis: None = None,
163+
out: None = None,
164+
overwrite_input: bool = False,
165+
method: _PercentileMethod = "linear",
166+
keepdims: _NoValueType | L[False] = ...,
167+
*,
168+
weights: CoFloating_1nd | None = None,
169+
interpolation: None = None,
170+
) -> Array[np.complexfloating]: ...
171+
@overload
172+
def nanpercentile(
173+
a: ToTimeDelta_nd,
174+
q: CoFloating_0d,
175+
axis: None = None,
176+
out: None = None,
177+
overwrite_input: bool = False,
178+
method: _PercentileMethod = "linear",
179+
keepdims: _NoValueType | L[False] = ...,
180+
*,
181+
weights: CoFloating_1nd | None = None,
182+
interpolation: None = None,
183+
) -> np.timedelta64: ...
184+
@overload
185+
def nanpercentile(
186+
a: ToTimeDelta_nd,
187+
q: CoFloating_1nd,
188+
axis: None = None,
189+
out: None = None,
190+
overwrite_input: bool = False,
191+
method: _PercentileMethod = "linear",
192+
keepdims: _NoValueType | L[False] = ...,
193+
*,
194+
weights: CoFloating_1nd | None = None,
195+
interpolation: None = None,
196+
) -> Array[np.timedelta64]: ...
197+
@overload
198+
def nanpercentile(
199+
a: ToDateTime_nd,
200+
q: CoFloating_0d,
201+
axis: None = None,
202+
out: None = None,
203+
overwrite_input: bool = False,
204+
method: _PercentileMethod = "linear",
205+
keepdims: _NoValueType | L[False] = ...,
206+
*,
207+
weights: CoFloating_1nd | None = None,
208+
interpolation: None = None,
209+
) -> np.datetime64: ...
210+
@overload
211+
def nanpercentile(
212+
a: ToDateTime_nd,
213+
q: CoFloating_1nd,
214+
axis: None = None,
215+
out: None = None,
216+
overwrite_input: bool = False,
217+
method: _PercentileMethod = "linear",
218+
keepdims: _NoValueType | L[False] = ...,
219+
*,
220+
weights: CoFloating_1nd | None = None,
221+
interpolation: None = None,
222+
) -> Array[np.datetime64]: ...
223+
@overload
224+
def nanpercentile(
225+
a: ToObject_nd,
226+
q: CoFloating_0d,
227+
axis: None = None,
228+
out: None = None,
229+
overwrite_input: bool = False,
230+
method: _PercentileMethod = "linear",
231+
keepdims: _NoValueType | L[False] = ...,
232+
*,
233+
weights: CoFloating_1nd | None = None,
234+
interpolation: None = None,
235+
) -> Any: ...
236+
@overload
237+
def nanpercentile(
238+
a: ToObject_nd,
239+
q: CoFloating_1nd,
240+
axis: None = None,
241+
out: None = None,
242+
overwrite_input: bool = False,
243+
method: _PercentileMethod = "linear",
244+
keepdims: _NoValueType | L[False] = ...,
245+
*,
246+
weights: CoFloating_1nd | None = None,
247+
interpolation: None = None,
248+
) -> Array[np.object_]: ...
249+
@overload
250+
def nanpercentile(
251+
a: CoComplex_nd | CoDateTime_nd | ToObject_nd,
252+
q: CoFloating_1nd,
253+
axis: _ShapeLike | None = None,
254+
out: None = None,
255+
overwrite_input: bool = False,
256+
method: _PercentileMethod = "linear",
257+
keepdims: _NoValueType | bool = ...,
258+
*,
259+
weights: CoFloating_1nd | None = None,
260+
interpolation: None = None,
261+
) -> Any: ...
262+
@overload
263+
def nanpercentile(
264+
a: CoComplex_nd | CoDateTime_nd | ToObject_nd,
265+
q: CoFloating_1nd,
266+
axis: _ShapeLike | None = None,
267+
*,
268+
out: _ArrayT,
269+
overwrite_input: bool = False,
270+
method: _PercentileMethod = "linear",
271+
keepdims: _NoValueType | bool = ...,
272+
weights: CoFloating_1nd | None = None,
273+
interpolation: None = None,
274+
) -> _ArrayT: ...
275+
@overload
276+
def nanpercentile(
277+
a: CoComplex_nd | CoDateTime_nd | ToObject_nd,
278+
q: CoFloating_1nd,
279+
axis: _ShapeLike | None,
280+
out: _ArrayT,
281+
overwrite_input: bool = False,
282+
method: _PercentileMethod = "linear",
283+
keepdims: _NoValueType | bool = ...,
284+
*,
285+
weights: CoFloating_1nd | None = None,
286+
interpolation: None = None,
287+
) -> _ArrayT: ...
288+
289+
nanquantile = nanpercentile

0 commit comments

Comments
 (0)