Skip to content

Commit 05f7019

Browse files
charrischarris
authored
Merge pull request numpy#27427 from jorenham/typing/ndarray-generic-all-any
TYP: Fixed & improved type-hinting for ``any`` and ``all``
2 parents 781e231 + a1b1f6e commit 05f7019

File tree

3 files changed

+185
-86
lines changed

3 files changed

+185
-86
lines changed

‎numpy/__init__.pyi

Lines changed: 137 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1448,62 +1448,6 @@ class _ArrayOrScalarCommon:
14481448
], /) -> None: ...
14491449
# an `np.bool` is returned when `keepdims=True` and `self` is a 0d array
14501450

1451-
@overload
1452-
def all(
1453-
self,
1454-
axis: None = ...,
1455-
out: None = ...,
1456-
keepdims: L[False] = ...,
1457-
*,
1458-
where: _ArrayLikeBool_co = ...,
1459-
) -> np.bool: ...
1460-
@overload
1461-
def all(
1462-
self,
1463-
axis: None | _ShapeLike = ...,
1464-
out: None = ...,
1465-
keepdims: builtins.bool = ...,
1466-
*,
1467-
where: _ArrayLikeBool_co = ...,
1468-
) -> Any: ...
1469-
@overload
1470-
def all(
1471-
self,
1472-
axis: None | _ShapeLike = ...,
1473-
out: _NdArraySubClass = ...,
1474-
keepdims: builtins.bool = ...,
1475-
*,
1476-
where: _ArrayLikeBool_co = ...,
1477-
) -> _NdArraySubClass: ...
1478-
1479-
@overload
1480-
def any(
1481-
self,
1482-
axis: None = ...,
1483-
out: None = ...,
1484-
keepdims: L[False] = ...,
1485-
*,
1486-
where: _ArrayLikeBool_co = ...,
1487-
) -> np.bool: ...
1488-
@overload
1489-
def any(
1490-
self,
1491-
axis: None | _ShapeLike = ...,
1492-
out: None = ...,
1493-
keepdims: builtins.bool = ...,
1494-
*,
1495-
where: _ArrayLikeBool_co = ...,
1496-
) -> Any: ...
1497-
@overload
1498-
def any(
1499-
self,
1500-
axis: None | _ShapeLike = ...,
1501-
out: _NdArraySubClass = ...,
1502-
keepdims: builtins.bool = ...,
1503-
*,
1504-
where: _ArrayLikeBool_co = ...,
1505-
) -> _NdArraySubClass: ...
1506-
15071451
@overload
15081452
def argmax(
15091453
self,
@@ -2027,6 +1971,80 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
20271971
@overload
20281972
def transpose(self, *axes: SupportsIndex) -> Self: ...
20291973

1974+
@overload
1975+
def all(
1976+
self,
1977+
axis: None = None,
1978+
out: None = None,
1979+
keepdims: L[False, 0] = False,
1980+
*,
1981+
where: _ArrayLikeBool_co = True
1982+
) -> np.bool: ...
1983+
@overload
1984+
def all(
1985+
self,
1986+
axis: None | int | tuple[int, ...] = None,
1987+
out: None = None,
1988+
keepdims: SupportsIndex = False,
1989+
*,
1990+
where: _ArrayLikeBool_co = True,
1991+
) -> np.bool | NDArray[np.bool]: ...
1992+
@overload
1993+
def all(
1994+
self,
1995+
axis: None | int | tuple[int, ...],
1996+
out: _NdArraySubClass,
1997+
keepdims: SupportsIndex = False,
1998+
*,
1999+
where: _ArrayLikeBool_co = True,
2000+
) -> _NdArraySubClass: ...
2001+
@overload
2002+
def all(
2003+
self,
2004+
axis: None | int | tuple[int, ...] = None,
2005+
*,
2006+
out: _NdArraySubClass,
2007+
keepdims: SupportsIndex = False,
2008+
where: _ArrayLikeBool_co = True,
2009+
) -> _NdArraySubClass: ...
2010+
2011+
@overload
2012+
def any(
2013+
self,
2014+
axis: None = None,
2015+
out: None = None,
2016+
keepdims: L[False, 0] = False,
2017+
*,
2018+
where: _ArrayLikeBool_co = True
2019+
) -> np.bool: ...
2020+
@overload
2021+
def any(
2022+
self,
2023+
axis: None | int | tuple[int, ...] = None,
2024+
out: None = None,
2025+
keepdims: SupportsIndex = False,
2026+
*,
2027+
where: _ArrayLikeBool_co = True,
2028+
) -> np.bool | NDArray[np.bool]: ...
2029+
@overload
2030+
def any(
2031+
self,
2032+
axis: None | int | tuple[int, ...],
2033+
out: _NdArraySubClass,
2034+
keepdims: SupportsIndex = False,
2035+
*,
2036+
where: _ArrayLikeBool_co = True,
2037+
) -> _NdArraySubClass: ...
2038+
@overload
2039+
def any(
2040+
self,
2041+
axis: None | int | tuple[int, ...] = None,
2042+
*,
2043+
out: _NdArraySubClass,
2044+
keepdims: SupportsIndex = False,
2045+
where: _ArrayLikeBool_co = True,
2046+
) -> _NdArraySubClass: ...
2047+
20302048
def argpartition(
20312049
self,
20322050
kth: _ArrayLikeInt_co,
@@ -3212,6 +3230,69 @@ class generic(_ArrayOrScalarCommon):
32123230

32133231
def squeeze(self, axis: None | L[0] | tuple[()] = ...) -> Self: ...
32143232
def transpose(self, axes: None | tuple[()] = ..., /) -> Self: ...
3233+
3234+
@overload
3235+
def all(
3236+
self,
3237+
/,
3238+
axis: L[0, -1] | tuple[()] | None = None,
3239+
out: None = None,
3240+
keepdims: SupportsIndex = False,
3241+
*,
3242+
where: builtins.bool | np.bool | ndarray[tuple[()], dtype[np.bool]] = True
3243+
) -> np.bool: ...
3244+
@overload
3245+
def all(
3246+
self,
3247+
/,
3248+
axis: L[0, -1] | tuple[()] | None,
3249+
out: ndarray[tuple[()], dtype[_SCT]],
3250+
keepdims: SupportsIndex = False,
3251+
*,
3252+
where: builtins.bool | np.bool | ndarray[tuple[()], dtype[np.bool]] = True,
3253+
) -> _SCT: ...
3254+
@overload
3255+
def all(
3256+
self,
3257+
/,
3258+
axis: L[0, -1] | tuple[()] | None = None,
3259+
*,
3260+
out: ndarray[tuple[()], dtype[_SCT]],
3261+
keepdims: SupportsIndex = False,
3262+
where: builtins.bool | np.bool | ndarray[tuple[()], dtype[np.bool]] = True,
3263+
) -> _SCT: ...
3264+
3265+
@overload
3266+
def any(
3267+
self,
3268+
/,
3269+
axis: L[0, -1] | tuple[()] | None = None,
3270+
out: None = None,
3271+
keepdims: SupportsIndex = False,
3272+
*,
3273+
where: builtins.bool | np.bool | ndarray[tuple[()], dtype[np.bool]] = True
3274+
) -> np.bool: ...
3275+
@overload
3276+
def any(
3277+
self,
3278+
/,
3279+
axis: L[0, -1] | tuple[()] | None,
3280+
out: ndarray[tuple[()], dtype[_SCT]],
3281+
keepdims: SupportsIndex = False,
3282+
*,
3283+
where: builtins.bool | np.bool | ndarray[tuple[()], dtype[np.bool]] = True,
3284+
) -> _SCT: ...
3285+
@overload
3286+
def any(
3287+
self,
3288+
/,
3289+
axis: L[0, -1] | tuple[()] | None = None,
3290+
*,
3291+
out: ndarray[tuple[()], dtype[_SCT]],
3292+
keepdims: SupportsIndex = False,
3293+
where: builtins.bool | np.bool | ndarray[tuple[()], dtype[np.bool]] = True,
3294+
) -> _SCT: ...
3295+
32153296
# Keep `dtype` at the bottom to avoid name conflicts with `np.dtype`
32163297
@property
32173298
def dtype(self) -> _dtype[Self]: ...

‎numpy/_core/fromnumeric.pyi

Lines changed: 44 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -575,57 +575,75 @@ def sum(
575575
@overload
576576
def all(
577577
a: ArrayLike,
578-
axis: None = ...,
579-
out: None = ...,
580-
keepdims: Literal[False] = ...,
578+
axis: None = None,
579+
out: None = None,
580+
keepdims: Literal[False, 0] = False,
581581
*,
582-
where: _ArrayLikeBool_co = ...,
582+
where: _ArrayLikeBool_co = True,
583583
) -> np.bool: ...
584584
@overload
585585
def all(
586586
a: ArrayLike,
587-
axis: None | _ShapeLike = ...,
588-
out: None = ...,
589-
keepdims: bool = ...,
587+
axis: None | int | tuple[int, ...] = None,
588+
out: None = None,
589+
keepdims: SupportsIndex = False,
590590
*,
591-
where: _ArrayLikeBool_co = ...,
592-
) -> Any: ...
591+
where: _ArrayLikeBool_co = True,
592+
) -> np.bool | NDArray[np.bool]: ...
593593
@overload
594594
def all(
595595
a: ArrayLike,
596-
axis: None | _ShapeLike = ...,
597-
out: _ArrayType = ...,
598-
keepdims: bool = ...,
596+
axis: None | int | tuple[int, ...],
597+
out: _ArrayType,
598+
keepdims: SupportsIndex = False,
599599
*,
600-
where: _ArrayLikeBool_co = ...,
600+
where: _ArrayLikeBool_co = True,
601+
) -> _ArrayType: ...
602+
@overload
603+
def all(
604+
a: ArrayLike,
605+
axis: None | int | tuple[int, ...] = None,
606+
*,
607+
out: _ArrayType,
608+
keepdims: SupportsIndex = False,
609+
where: _ArrayLikeBool_co = True,
601610
) -> _ArrayType: ...
602611

603612
@overload
604613
def any(
605614
a: ArrayLike,
606-
axis: None = ...,
607-
out: None = ...,
608-
keepdims: Literal[False] = ...,
615+
axis: None = None,
616+
out: None = None,
617+
keepdims: Literal[False, 0] = False,
609618
*,
610-
where: _ArrayLikeBool_co = ...,
619+
where: _ArrayLikeBool_co = True,
611620
) -> np.bool: ...
612621
@overload
613622
def any(
614623
a: ArrayLike,
615-
axis: None | _ShapeLike = ...,
616-
out: None = ...,
617-
keepdims: bool = ...,
624+
axis: None | int | tuple[int, ...] = None,
625+
out: None = None,
626+
keepdims: SupportsIndex = False,
618627
*,
619-
where: _ArrayLikeBool_co = ...,
620-
) -> Any: ...
628+
where: _ArrayLikeBool_co = True,
629+
) -> np.bool | NDArray[np.bool]: ...
621630
@overload
622631
def any(
623632
a: ArrayLike,
624-
axis: None | _ShapeLike = ...,
625-
out: _ArrayType = ...,
626-
keepdims: bool = ...,
633+
axis: None | int | tuple[int, ...],
634+
out: _ArrayType,
635+
keepdims: SupportsIndex = False,
627636
*,
628-
where: _ArrayLikeBool_co = ...,
637+
where: _ArrayLikeBool_co = True,
638+
) -> _ArrayType: ...
639+
@overload
640+
def any(
641+
a: ArrayLike,
642+
axis: None | int | tuple[int, ...] = None,
643+
*,
644+
out: _ArrayType,
645+
keepdims: SupportsIndex = False,
646+
where: _ArrayLikeBool_co = True,
629647
) -> _ArrayType: ...
630648

631649
@overload

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,14 +44,14 @@ assert_type(ctypes_obj.strides_as(ct.c_ubyte), ct.Array[ct.c_ubyte])
4444

4545
assert_type(f8.all(), np.bool)
4646
assert_type(AR_f8.all(), np.bool)
47-
assert_type(AR_f8.all(axis=0), Any)
48-
assert_type(AR_f8.all(keepdims=True), Any)
47+
assert_type(AR_f8.all(axis=0), np.bool | npt.NDArray[np.bool])
48+
assert_type(AR_f8.all(keepdims=True), np.bool | npt.NDArray[np.bool])
4949
assert_type(AR_f8.all(out=B), SubClass)
5050

5151
assert_type(f8.any(), np.bool)
5252
assert_type(AR_f8.any(), np.bool)
53-
assert_type(AR_f8.any(axis=0), Any)
54-
assert_type(AR_f8.any(keepdims=True), Any)
53+
assert_type(AR_f8.any(axis=0), np.bool | npt.NDArray[np.bool])
54+
assert_type(AR_f8.any(keepdims=True), np.bool | npt.NDArray[np.bool])
5555
assert_type(AR_f8.any(out=B), SubClass)
5656

5757
assert_type(f8.argmax(), np.intp)

0 commit comments

Comments
 (0)