Skip to content

Commit 35b02ec

Browse files
guan404mingjorenham
guan404ming
and
jorenham
authored
🏷️ ufunc annotations for logical_{not,and,or,xor} (#324)
* 🏷️ ufunc annotations for logical-{not,and,or,xor} * ♻️ refactor ufunc annotations for logical operations * ♻️ refactor annotation for `logical` related class Co-Authored-By: Joren Hammudoglu <jhammudoglu@gmail.com> * ♻️ refactor ufunc annotations for logical operations * 🏷️ fix keyword annotation Co-Authored-By: Joren Hammudoglu <jhammudoglu@gmail.com> * 🔥 reorder keyword para * 🐛 update logical type Co-Authored-By: Joren Hammudoglu <jhammudoglu@gmail.com> * 🐛 update src/numpy-stubs/_typing/_ufunc.pyi Co-authored-by: Joren Hammudoglu <jhammudoglu@gmail.com> --------- Co-authored-by: Joren Hammudoglu <jhammudoglu@gmail.com>
1 parent 20cdeb7 commit 35b02ec

File tree

4 files changed

+252
-8
lines changed

4 files changed

+252
-8
lines changed

‎src/numpy-stubs/__init__.pyi‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -309,7 +309,18 @@ from ._typing._char_codes import (
309309
_UnsignedIntegerCodes,
310310
_VoidCodes,
311311
)
312-
from ._typing._ufunc import _Call11Bool, _Call11Isnat, _Call21Bool, _gufunc_2_1, _ufunc_1_1, _ufunc_1_2, _ufunc_2_1, _ufunc_2_2
312+
from ._typing._ufunc import (
313+
_Call11Bool,
314+
_Call11Isnat,
315+
_Call11Logical,
316+
_Call21Bool,
317+
_Call21Logical,
318+
_gufunc_2_1,
319+
_ufunc_1_1,
320+
_ufunc_1_2,
321+
_ufunc_2_1,
322+
_ufunc_2_2,
323+
)
313324
from .lib import scimath as emath
314325
from .lib._arraypad_impl import pad
315326
from .lib._arraysetops_impl import (
@@ -7103,7 +7114,7 @@ log: Final[_ufunc_1_1] = ...
71037114
log2: Final[_ufunc_1_1] = ...
71047115
log10: Final[_ufunc_1_1] = ...
71057116
log1p: Final[_ufunc_1_1] = ...
7106-
logical_not: Final[_ufunc_1_1] = ...
7117+
logical_not: Final[_ufunc_1_1[_Call11Logical]] = ...
71077118
negative: Final[_ufunc_1_1] = ...
71087119
positive: Final[_ufunc_1_1] = ...
71097120
rad2deg: Final[_ufunc_1_1] = ...
@@ -7160,9 +7171,9 @@ ldexp: Final[_ufunc_2_1] = ...
71607171
left_shift: Final[_ufunc_2_1] = ...
71617172
logaddexp: Final[_ufunc_2_1] = ...
71627173
logaddexp2: Final[_ufunc_2_1] = ...
7163-
logical_and: Final[_ufunc_2_1] = ...
7164-
logical_or: Final[_ufunc_2_1] = ...
7165-
logical_xor: Final[_ufunc_2_1] = ...
7174+
logical_and: Final[_ufunc_2_1[_Call21Logical]] = ...
7175+
logical_or: Final[_ufunc_2_1[_Call21Logical]] = ...
7176+
logical_xor: Final[_ufunc_2_1[_Call21Logical]] = ...
71667177
maximum: Final[_ufunc_2_1] = ...
71677178
minimum: Final[_ufunc_2_1] = ...
71687179
mod: Final[_ufunc_2_1] = ...

‎src/numpy-stubs/_typing/_ufunc.pyi‎

Lines changed: 197 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,17 @@ import numpy as np
1616
from numpy import _CastingKind, _OrderKACF # noqa: ICN003
1717
from numpy._typing import _DTypeLikeBool, _NestedSequence
1818

19-
from ._array_like import ArrayLike, NDArray, _ArrayLike, _ArrayLikeBool_co, _ArrayLikeInt_co
19+
from ._array_like import (
20+
ArrayLike,
21+
NDArray,
22+
_ArrayLike,
23+
_ArrayLikeBool_co,
24+
_ArrayLikeInt_co,
25+
_ArrayLikeNumber_co,
26+
_ArrayLikeObject_co,
27+
)
2028
from ._dtype_like import DTypeLike, _DTypeLike
21-
from ._scalars import _ScalarLike_co
29+
from ._scalars import _NumberLike_co, _ScalarLike_co
2230
from ._shape import _ShapeLike
2331

2432
###
@@ -201,7 +209,7 @@ class _Call11Bool(Protocol):
201209
dtype: _DTypeLikeBool | None = None,
202210
**kwds: Unpack[_Kwargs2],
203211
) -> _ArrayT: ...
204-
@overload # (array) -> Array[bool] | bool
212+
@overload # (array) -> Array[bool]
205213
def __call__(
206214
self,
207215
x: _AnyArray,
@@ -255,6 +263,79 @@ class _Call11Isnat(Protocol):
255263
**kwds: Unpack[_Kwargs2],
256264
) -> NDArray[np.bool]: ...
257265

266+
@type_check_only
267+
class _Call11Logical(Protocol):
268+
@overload
269+
def __call__( # (scalar, dtype: np.object_) -> bool
270+
self,
271+
x: _ScalarLike_co,
272+
/,
273+
out: None = None,
274+
*,
275+
dtype: _DTypeLike[np.object_],
276+
**kwargs: Unpack[_Kwargs2],
277+
) -> bool: ...
278+
@overload
279+
def __call__( # (scalar) -> np.bool
280+
self,
281+
x: _NumberLike_co,
282+
/,
283+
out: None = None,
284+
*,
285+
dtype: _DTypeLikeBool | None = None,
286+
**kwargs: Unpack[_Kwargs2],
287+
) -> np.bool: ...
288+
@overload
289+
def __call__( # (array-like, dtype: np.object_) -> np.object_
290+
self,
291+
x: _ArrayLikeNumber_co | _ArrayLikeObject_co,
292+
/,
293+
out: None = None,
294+
*,
295+
dtype: _DTypeLike[np.object_],
296+
**kwargs: Unpack[_Kwargs2],
297+
) -> NDArray[np.object_] | bool: ...
298+
@overload
299+
def __call__( # (array-like, out: T) -> T
300+
self,
301+
x: _ArrayLikeNumber_co | _ArrayLikeObject_co,
302+
/,
303+
out: _Out1[_ArrayT],
304+
*,
305+
dtype: DTypeLike | None = None,
306+
**kwargs: Unpack[_Kwargs2],
307+
) -> _ArrayT: ...
308+
@overload # (array) -> Array[bool]
309+
def __call__(
310+
self,
311+
x: NDArray[np.bool | np.number] | _NestedSequence[np.bool | np.number],
312+
/,
313+
out: _Out1[NDArray[np.bool]] | None = None,
314+
*,
315+
dtype: _DTypeLikeBool | None = None,
316+
**kwds: Unpack[_Kwargs2],
317+
) -> NDArray[np.bool]: ...
318+
@overload
319+
def __call__( # (array-like) -> Array[bool] | bool
320+
self,
321+
x: _ArrayLikeNumber_co,
322+
/,
323+
out: None = None,
324+
*,
325+
dtype: _DTypeLikeBool | None = None,
326+
**kwargs: Unpack[_Kwargs2],
327+
) -> NDArray[np.bool] | np.bool: ...
328+
@overload
329+
def __call__( # (?) -> ?
330+
self,
331+
x: _CanArrayUFunc,
332+
/,
333+
out: _Out1[_AnyArray] | None = None,
334+
*,
335+
dtype: DTypeLike | None = None,
336+
**kwargs: Unpack[_Kwargs2],
337+
) -> Any: ...
338+
258339
@type_check_only
259340
class _Call12(Protocol):
260341
@overload
@@ -418,6 +499,119 @@ class _Call21Bool(Protocol):
418499
**kwds: Unpack[_Kwargs3],
419500
) -> np.bool | NDArray[np.bool]: ...
420501

502+
@type_check_only
503+
class _Call21Logical(Protocol):
504+
@overload # (scalar, scalar, dtype: np.object_) -> np.object_
505+
def __call__(
506+
self,
507+
x1: _ScalarLike_co,
508+
x2: _ScalarLike_co,
509+
/,
510+
out: None = None,
511+
*,
512+
dtype: _DTypeLike[np.object_],
513+
**kwds: Unpack[_Kwargs3],
514+
) -> bool: ...
515+
@overload # (scalar, scalar) -> bool
516+
def __call__(
517+
self,
518+
x1: _NumberLike_co,
519+
x2: _NumberLike_co,
520+
/,
521+
out: None = None,
522+
*,
523+
dtype: _DTypeLikeBool | None = None,
524+
**kwds: Unpack[_Kwargs3],
525+
) -> np.bool: ...
526+
@overload # (array-like, array, dtype: object_) -> Array[object_]
527+
def __call__(
528+
self,
529+
x1: _ArrayLikeNumber_co | _ArrayLikeObject_co,
530+
x2: _AnyArray,
531+
/,
532+
out: None = None,
533+
*,
534+
dtype: _DTypeLike[np.object_],
535+
**kwds: Unpack[_Kwargs3],
536+
) -> NDArray[np.object_]: ...
537+
@overload # (array, array-like, dtype: object_) -> Array[object_]
538+
def __call__(
539+
self,
540+
x1: _AnyArray,
541+
x2: _ArrayLikeNumber_co | _ArrayLikeObject_co,
542+
/,
543+
out: None = None,
544+
*,
545+
dtype: _DTypeLike[np.object_],
546+
**kwds: Unpack[_Kwargs3],
547+
) -> NDArray[np.object_]: ...
548+
@overload # (array-like, array, dtype: dtype[T]) -> Array[T]
549+
def __call__(
550+
self,
551+
x1: _ArrayLikeNumber_co,
552+
x2: NDArray[np.bool | np.number] | _NestedSequence[np.bool | np.number],
553+
/,
554+
out: None = None,
555+
*,
556+
dtype: _DTypeLikeBool | None = None,
557+
**kwds: Unpack[_Kwargs3],
558+
) -> NDArray[np.bool]: ...
559+
@overload # (array, array-like, dtype: dtype[T]) -> Array[T]
560+
def __call__(
561+
self,
562+
x1: NDArray[np.bool | np.number] | _NestedSequence[np.bool | np.number],
563+
x2: _ArrayLikeNumber_co,
564+
/,
565+
out: None = None,
566+
*,
567+
dtype: _DTypeLikeBool | None = None,
568+
**kwds: Unpack[_Kwargs3],
569+
) -> NDArray[np.bool]: ...
570+
@overload # (array-like, array-like, out: T) -> T
571+
def __call__(
572+
self,
573+
x1: _ArrayLikeNumber_co | _ArrayLikeObject_co,
574+
x2: _ArrayLikeNumber_co | _ArrayLikeObject_co,
575+
/,
576+
out: _Out1[_ArrayT],
577+
*,
578+
dtype: None = None,
579+
**kwds: Unpack[_Kwargs3],
580+
) -> _ArrayT: ...
581+
@overload # (array-like, array) -> Array[?]
582+
def __call__(
583+
self,
584+
x1: _ArrayLikeNumber_co,
585+
x2: NDArray[np.bool | np.number] | _NestedSequence[np.bool | np.number | complex],
586+
/,
587+
out: _Out1[NDArray[np.bool]] | None = None,
588+
*,
589+
dtype: _DTypeLikeBool | None = None,
590+
**kwds: Unpack[_Kwargs3],
591+
) -> NDArray[np.bool]: ...
592+
@overload # (array, array-like) -> Array[?]
593+
def __call__(
594+
self,
595+
x1: NDArray[np.bool | np.number] | _NestedSequence[np.bool | np.number | complex],
596+
x2: _ArrayLikeNumber_co,
597+
/,
598+
out: _Out1[NDArray[np.bool]] | None = None,
599+
*,
600+
dtype: _DTypeLikeBool | None = None,
601+
**kwds: Unpack[_Kwargs3],
602+
) -> NDArray[np.bool]: ...
603+
@overload # (array-like, array-like) -> Array[?] | ?
604+
def __call__(
605+
self,
606+
x1: _ArrayLikeNumber_co | _ArrayLikeObject_co,
607+
x2: _ArrayLikeNumber_co | _ArrayLikeObject_co,
608+
/,
609+
out: _Out1[_AnyArray] | None = None,
610+
*,
611+
dtype: DTypeLike | None = None,
612+
**kwds: Unpack[_Kwargs3],
613+
) -> Any: ...
614+
421615
@type_check_only
422616
class _Call21(Protocol):
423617
@overload # (scalar, scalar, dtype: type[T]) -> T

‎test/static/accept/ufuncs.pyi‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,3 +62,29 @@ assert_type(np.isinf(AR_f8, out=AR_bool), npt.NDArray[np.bool_])
6262
assert_type(np.isfinite(f8), np.bool_)
6363
assert_type(np.isfinite(AR_f8), npt.NDArray[np.bool_])
6464
assert_type(np.isfinite(AR_f8, out=AR_bool), npt.NDArray[np.bool_])
65+
66+
assert_type(np.logical_not(True), np.bool_)
67+
assert_type(np.logical_not(AR_bool), npt.NDArray[np.bool_])
68+
assert_type(np.logical_not(AR_bool, out=AR_bool), npt.NDArray[np.bool_])
69+
assert_type(np.logical_not(AR_bool, dtype=np.object_), npt.NDArray[np.object_] | bool)
70+
71+
assert_type(np.logical_and(True, True), np.bool_)
72+
assert_type(np.logical_and(AR_bool, AR_bool), npt.NDArray[np.bool_])
73+
assert_type(np.logical_and(AR_bool, AR_bool, out=AR_bool), npt.NDArray[np.bool_])
74+
assert_type(np.logical_and(AR_i8, AR_bool, out=AR_bool), npt.NDArray[np.bool_])
75+
assert_type(np.logical_and(AR_bool, AR_i8), npt.NDArray[np.bool_])
76+
assert_type(np.logical_and(AR_bool, AR_bool, dtype=np.object_), npt.NDArray[np.object_])
77+
78+
assert_type(np.logical_or(True, True), np.bool_)
79+
assert_type(np.logical_or(AR_bool, AR_bool), npt.NDArray[np.bool_])
80+
assert_type(np.logical_or(AR_bool, AR_bool, out=AR_bool), npt.NDArray[np.bool_])
81+
assert_type(np.logical_or(AR_i8, AR_bool, out=AR_bool), npt.NDArray[np.bool_])
82+
assert_type(np.logical_or(AR_bool, AR_i8), npt.NDArray[np.bool_])
83+
assert_type(np.logical_or(AR_bool, AR_bool, dtype=np.object_), npt.NDArray[np.object_])
84+
85+
assert_type(np.logical_xor(True, True), np.bool_)
86+
assert_type(np.logical_xor(AR_bool, AR_bool), npt.NDArray[np.bool_])
87+
assert_type(np.logical_xor(AR_bool, AR_bool, out=AR_bool), npt.NDArray[np.bool_])
88+
assert_type(np.logical_xor(AR_i8, AR_bool, out=AR_bool), npt.NDArray[np.bool_])
89+
assert_type(np.logical_xor(AR_bool, AR_i8), npt.NDArray[np.bool_])
90+
assert_type(np.logical_xor(AR_bool, AR_bool, dtype=np.object_), npt.NDArray[np.object_])

‎test/static/reject/ufuncs.pyi‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import numpy.typing as npt
33

44
i8: np.int64
55
AR_f8: npt.NDArray[np.float64]
6+
dt64: np.datetime64
67

78
np.sin.nin + "foo" # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
89

@@ -47,3 +48,15 @@ np.isnat(i8, dtype=np.int64) # type: ignore[call-overload] # pyright: ignore[r
4748
np.isnat(i8) # type: ignore[call-overload] # pyright: ignore[reportCallIssue,reportArgumentType]
4849
np.isinf(i8, dtype=np.int64) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
4950
np.isfinite(i8, dtype=np.int64) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
51+
52+
np.logical_not(i8, dtype=np.datetime64) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
53+
np.logical_not(dt64) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
54+
55+
np.logical_and(dt64, dt64) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
56+
np.logical_and(dt64, dt64, dtype=np.datetime64) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
57+
58+
np.logical_or(dt64, dt64) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
59+
np.logical_or(dt64, dt64, dtype=np.datetime64) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
60+
61+
np.logical_xor(dt64, dt64) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
62+
np.logical_xor(dt64, dt64, dtype=np.datetime64) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]

0 commit comments

Comments
 (0)