Skip to content

Commit 51fa653

Browse files
rgommersrgommers
authored
Merge pull request numpy#27100 from jorenham/typing/27092-histogram2d-bins-float
TYP: Fixed & improved type hints for ``numpy.histogram2d``
2 parents b664fe0 + 636c275 commit 51fa653

File tree

2 files changed

+264
-28
lines changed

2 files changed

+264
-28
lines changed

‎numpy/lib/_twodim_base_impl.pyi

Lines changed: 200 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import builtins
22
from collections.abc import Callable, Sequence
33
from typing import (
44
Any,
5+
TypeAlias,
56
overload,
67
TypeVar,
78
Literal as L,
@@ -16,6 +17,7 @@ from numpy import (
1617
int_,
1718
intp,
1819
float64,
20+
complex128,
1921
signedinteger,
2022
floating,
2123
complexfloating,
@@ -29,6 +31,7 @@ from numpy._typing import (
2931
ArrayLike,
3032
_ArrayLike,
3133
NDArray,
34+
_SupportsArray,
3235
_SupportsArrayFunc,
3336
_ArrayLikeInt_co,
3437
_ArrayLikeFloat_co,
@@ -164,44 +167,220 @@ def vander(
164167
increasing: bool = ...,
165168
) -> NDArray[object_]: ...
166169

170+
171+
_Int_co: TypeAlias = np.integer[Any] | np.bool
172+
_Float_co: TypeAlias = np.floating[Any] | _Int_co
173+
_Number_co: TypeAlias = np.number[Any] | np.bool
174+
175+
_ArrayLike1D: TypeAlias = _SupportsArray[np.dtype[_SCT]] | Sequence[_SCT]
176+
_ArrayLike2D: TypeAlias = (
177+
_SupportsArray[np.dtype[_SCT]]
178+
| Sequence[_ArrayLike1D[_SCT]]
179+
)
180+
181+
_ArrayLike1DInt_co = (
182+
_SupportsArray[np.dtype[_Int_co]]
183+
| Sequence[int | _Int_co]
184+
)
185+
_ArrayLike1DFloat_co = (
186+
_SupportsArray[np.dtype[_Float_co]]
187+
| Sequence[float | int | _Float_co]
188+
)
189+
_ArrayLike2DFloat_co = (
190+
_SupportsArray[np.dtype[_Float_co]]
191+
| Sequence[_ArrayLike1DFloat_co]
192+
)
193+
_ArrayLike1DNumber_co = (
194+
_SupportsArray[np.dtype[_Number_co]]
195+
| Sequence[int | float | complex | _Number_co]
196+
)
197+
198+
_SCT_complex = TypeVar("_SCT_complex", bound=np.complexfloating[Any, Any])
199+
_SCT_inexact = TypeVar("_SCT_inexact", bound=np.inexact[Any])
200+
_SCT_number_co = TypeVar("_SCT_number_co", bound=_Number_co)
201+
167202
@overload
168-
def histogram2d( # type: ignore[misc]
169-
x: _ArrayLikeFloat_co,
170-
y: _ArrayLikeFloat_co,
203+
def histogram2d(
204+
x: _ArrayLike1D[_SCT_complex],
205+
y: _ArrayLike1D[_SCT_complex | _Float_co],
171206
bins: int | Sequence[int] = ...,
172-
range: None | _ArrayLikeFloat_co = ...,
207+
range: None | _ArrayLike2DFloat_co = ...,
173208
density: None | bool = ...,
174-
weights: None | _ArrayLikeFloat_co = ...,
209+
weights: None | _ArrayLike1DFloat_co = ...,
175210
) -> tuple[
176211
NDArray[float64],
177-
NDArray[floating[Any]],
178-
NDArray[floating[Any]],
212+
NDArray[_SCT_complex],
213+
NDArray[_SCT_complex],
179214
]: ...
180215
@overload
181216
def histogram2d(
182-
x: _ArrayLikeComplex_co,
183-
y: _ArrayLikeComplex_co,
217+
x: _ArrayLike1D[_SCT_complex | _Float_co],
218+
y: _ArrayLike1D[_SCT_complex],
184219
bins: int | Sequence[int] = ...,
185-
range: None | _ArrayLikeFloat_co = ...,
220+
range: None | _ArrayLike2DFloat_co = ...,
186221
density: None | bool = ...,
187-
weights: None | _ArrayLikeFloat_co = ...,
222+
weights: None | _ArrayLike1DFloat_co = ...,
188223
) -> tuple[
189224
NDArray[float64],
190-
NDArray[complexfloating[Any, Any]],
191-
NDArray[complexfloating[Any, Any]],
225+
NDArray[_SCT_complex],
226+
NDArray[_SCT_complex],
192227
]: ...
193-
@overload # TODO: Sort out `bins`
228+
@overload
194229
def histogram2d(
195-
x: _ArrayLikeComplex_co,
196-
y: _ArrayLikeComplex_co,
197-
bins: Sequence[_ArrayLikeInt_co],
198-
range: None | _ArrayLikeFloat_co = ...,
230+
x: _ArrayLike1D[_SCT_inexact],
231+
y: _ArrayLike1D[_SCT_inexact | _Int_co],
232+
bins: int | Sequence[int] = ...,
233+
range: None | _ArrayLike2DFloat_co = ...,
234+
density: None | bool = ...,
235+
weights: None | _ArrayLike1DFloat_co = ...,
236+
) -> tuple[
237+
NDArray[float64],
238+
NDArray[_SCT_inexact],
239+
NDArray[_SCT_inexact],
240+
]: ...
241+
@overload
242+
def histogram2d(
243+
x: _ArrayLike1D[_SCT_inexact | _Int_co],
244+
y: _ArrayLike1D[_SCT_inexact],
245+
bins: int | Sequence[int] = ...,
246+
range: None | _ArrayLike2DFloat_co = ...,
247+
density: None | bool = ...,
248+
weights: None | _ArrayLike1DFloat_co = ...,
249+
) -> tuple[
250+
NDArray[float64],
251+
NDArray[_SCT_inexact],
252+
NDArray[_SCT_inexact],
253+
]: ...
254+
@overload
255+
def histogram2d(
256+
x: _ArrayLike1DInt_co | Sequence[float | int],
257+
y: _ArrayLike1DInt_co | Sequence[float | int],
258+
bins: int | Sequence[int] = ...,
259+
range: None | _ArrayLike2DFloat_co = ...,
260+
density: None | bool = ...,
261+
weights: None | _ArrayLike1DFloat_co = ...,
262+
) -> tuple[
263+
NDArray[float64],
264+
NDArray[float64],
265+
NDArray[float64],
266+
]: ...
267+
@overload
268+
def histogram2d(
269+
x: Sequence[complex | float | int],
270+
y: Sequence[complex | float | int],
271+
bins: int | Sequence[int] = ...,
272+
range: None | _ArrayLike2DFloat_co = ...,
273+
density: None | bool = ...,
274+
weights: None | _ArrayLike1DFloat_co = ...,
275+
) -> tuple[
276+
NDArray[float64],
277+
NDArray[complex128 | float64],
278+
NDArray[complex128 | float64],
279+
]: ...
280+
@overload
281+
def histogram2d(
282+
x: _ArrayLike1DNumber_co,
283+
y: _ArrayLike1DNumber_co,
284+
bins: _ArrayLike1D[_SCT_number_co] | Sequence[_ArrayLike1D[_SCT_number_co]],
285+
range: None | _ArrayLike2DFloat_co = ...,
286+
density: None | bool = ...,
287+
weights: None | _ArrayLike1DFloat_co = ...,
288+
) -> tuple[
289+
NDArray[float64],
290+
NDArray[_SCT_number_co],
291+
NDArray[_SCT_number_co],
292+
]: ...
293+
@overload
294+
def histogram2d(
295+
x: _ArrayLike1D[_SCT_inexact],
296+
y: _ArrayLike1D[_SCT_inexact],
297+
bins: Sequence[_ArrayLike1D[_SCT_number_co] | int],
298+
range: None | _ArrayLike2DFloat_co = ...,
299+
density: None | bool = ...,
300+
weights: None | _ArrayLike1DFloat_co = ...,
301+
) -> tuple[
302+
NDArray[float64],
303+
NDArray[_SCT_number_co | _SCT_inexact],
304+
NDArray[_SCT_number_co | _SCT_inexact],
305+
]: ...
306+
@overload
307+
def histogram2d(
308+
x: _ArrayLike1DInt_co | Sequence[float | int],
309+
y: _ArrayLike1DInt_co | Sequence[float | int],
310+
bins: Sequence[_ArrayLike1D[_SCT_number_co] | int],
311+
range: None | _ArrayLike2DFloat_co = ...,
312+
density: None | bool = ...,
313+
weights: None | _ArrayLike1DFloat_co = ...,
314+
) -> tuple[
315+
NDArray[float64],
316+
NDArray[_SCT_number_co | float64],
317+
NDArray[_SCT_number_co | float64],
318+
]: ...
319+
@overload
320+
def histogram2d(
321+
x: Sequence[complex | float | int],
322+
y: Sequence[complex | float | int],
323+
bins: Sequence[_ArrayLike1D[_SCT_number_co] | int],
324+
range: None | _ArrayLike2DFloat_co = ...,
325+
density: None | bool = ...,
326+
weights: None | _ArrayLike1DFloat_co = ...,
327+
) -> tuple[
328+
NDArray[float64],
329+
NDArray[_SCT_number_co | complex128 | float64],
330+
NDArray[_SCT_number_co | complex128 | float64] ,
331+
]: ...
332+
333+
@overload
334+
def histogram2d(
335+
x: _ArrayLike1DNumber_co,
336+
y: _ArrayLike1DNumber_co,
337+
bins: Sequence[Sequence[bool]],
338+
range: None | _ArrayLike2DFloat_co = ...,
339+
density: None | bool = ...,
340+
weights: None | _ArrayLike1DFloat_co = ...,
341+
) -> tuple[
342+
NDArray[float64],
343+
NDArray[np.bool],
344+
NDArray[np.bool],
345+
]: ...
346+
@overload
347+
def histogram2d(
348+
x: _ArrayLike1DNumber_co,
349+
y: _ArrayLike1DNumber_co,
350+
bins: Sequence[Sequence[int | bool]],
351+
range: None | _ArrayLike2DFloat_co = ...,
352+
density: None | bool = ...,
353+
weights: None | _ArrayLike1DFloat_co = ...,
354+
) -> tuple[
355+
NDArray[float64],
356+
NDArray[np.int_ | np.bool],
357+
NDArray[np.int_ | np.bool],
358+
]: ...
359+
@overload
360+
def histogram2d(
361+
x: _ArrayLike1DNumber_co,
362+
y: _ArrayLike1DNumber_co,
363+
bins: Sequence[Sequence[float | int | bool]],
364+
range: None | _ArrayLike2DFloat_co = ...,
365+
density: None | bool = ...,
366+
weights: None | _ArrayLike1DFloat_co = ...,
367+
) -> tuple[
368+
NDArray[float64],
369+
NDArray[np.float64 | np.int_ | np.bool],
370+
NDArray[np.float64 | np.int_ | np.bool],
371+
]: ...
372+
@overload
373+
def histogram2d(
374+
x: _ArrayLike1DNumber_co,
375+
y: _ArrayLike1DNumber_co,
376+
bins: Sequence[Sequence[complex | float | int | bool]],
377+
range: None | _ArrayLike2DFloat_co = ...,
199378
density: None | bool = ...,
200-
weights: None | _ArrayLikeFloat_co = ...,
379+
weights: None | _ArrayLike1DFloat_co = ...,
201380
) -> tuple[
202381
NDArray[float64],
203-
NDArray[Any],
204-
NDArray[Any],
382+
NDArray[np.complex128 | np.float64 | np.int_ | np.bool],
383+
NDArray[np.complex128 | np.float64 | np.int_ | np.bool],
205384
]: ...
206385

207386
# NOTE: we're assuming/demanding here the `mask_func` returns

‎numpy/typing/tests/data/reveal/twodim_base.pyi

Lines changed: 64 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ AR_c: npt.NDArray[np.complex128]
2828
AR_O: npt.NDArray[np.object_]
2929

3030
AR_LIKE_b: list[bool]
31+
AR_LIKE_c: list[complex]
3132

3233
assert_type(np.fliplr(AR_b), npt.NDArray[np.bool])
3334
assert_type(np.fliplr(AR_LIKE_b), npt.NDArray[Any])
@@ -62,28 +63,84 @@ assert_type(np.vander(AR_f, increasing=True), npt.NDArray[np.floating[Any]])
6263
assert_type(np.vander(AR_c), npt.NDArray[np.complexfloating[Any, Any]])
6364
assert_type(np.vander(AR_O), npt.NDArray[np.object_])
6465

66+
assert_type(
67+
np.histogram2d(AR_LIKE_c, AR_LIKE_c),
68+
tuple[
69+
npt.NDArray[np.float64],
70+
npt.NDArray[np.complex128 | np.float64],
71+
npt.NDArray[np.complex128 | np.float64],
72+
],
73+
)
6574
assert_type(
6675
np.histogram2d(AR_i, AR_b),
6776
tuple[
6877
npt.NDArray[np.float64],
69-
npt.NDArray[np.floating[Any]],
70-
npt.NDArray[np.floating[Any]],
78+
npt.NDArray[np.float64],
79+
npt.NDArray[np.float64],
7180
],
7281
)
7382
assert_type(
74-
np.histogram2d(AR_f, AR_f),
83+
np.histogram2d(AR_f, AR_i),
7584
tuple[
7685
npt.NDArray[np.float64],
77-
npt.NDArray[np.floating[Any]],
78-
npt.NDArray[np.floating[Any]],
86+
npt.NDArray[np.float64],
87+
npt.NDArray[np.float64],
88+
],
89+
)
90+
assert_type(
91+
np.histogram2d(AR_i, AR_f),
92+
tuple[
93+
npt.NDArray[np.float64],
94+
npt.NDArray[np.float64],
95+
npt.NDArray[np.float64],
7996
],
8097
)
8198
assert_type(
8299
np.histogram2d(AR_f, AR_c, weights=AR_LIKE_b),
83100
tuple[
84101
npt.NDArray[np.float64],
85-
npt.NDArray[np.complexfloating[Any, Any]],
86-
npt.NDArray[np.complexfloating[Any, Any]],
102+
npt.NDArray[np.complex128],
103+
npt.NDArray[np.complex128],
104+
],
105+
)
106+
assert_type(
107+
np.histogram2d(AR_f, AR_c, bins=8),
108+
tuple[
109+
npt.NDArray[np.float64],
110+
npt.NDArray[np.complex128],
111+
npt.NDArray[np.complex128],
112+
],
113+
)
114+
assert_type(
115+
np.histogram2d(AR_c, AR_f, bins=(8, 5)),
116+
tuple[
117+
npt.NDArray[np.float64],
118+
npt.NDArray[np.complex128],
119+
npt.NDArray[np.complex128],
120+
],
121+
)
122+
assert_type(
123+
np.histogram2d(AR_c, AR_i, bins=AR_u),
124+
tuple[
125+
npt.NDArray[np.float64],
126+
npt.NDArray[np.uint64],
127+
npt.NDArray[np.uint64],
128+
],
129+
)
130+
assert_type(
131+
np.histogram2d(AR_c, AR_c, bins=(AR_u, AR_u)),
132+
tuple[
133+
npt.NDArray[np.float64],
134+
npt.NDArray[np.uint64],
135+
npt.NDArray[np.uint64],
136+
],
137+
)
138+
assert_type(
139+
np.histogram2d(AR_c, AR_c, bins=(AR_b, 8)),
140+
tuple[
141+
npt.NDArray[np.float64],
142+
npt.NDArray[np.bool | np.complex128],
143+
npt.NDArray[np.bool | np.complex128],
87144
],
88145
)
89146

0 commit comments

Comments
 (0)