8000 ENH: Add annotations for bitwise operations by BvB93 · Pull Request #17465 · numpy/numpy · GitHub
[go: up one dir, main page]

Skip to content

ENH: Add annotations for bitwise operations #17465

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Oct 9, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 63 additions & 16 deletions numpy/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,15 @@ from numpy.core._internal import _ctypes
from numpy.typing import ArrayLike, DtypeLike, _Shape, _ShapeLike
from numpy.typing._callable import (
_BoolOp,
_BoolBitOp,
_BoolSub,
_BoolTrueDiv,
_TD64Div,
_IntTrueDiv,
_UnsignedIntOp,
_UnsignedIntBitOp,
_SignedIntOp,
_SignedIntBitOp,
_FloatOp,
_ComplexOp,
_NumberOp,
Expand Down Expand Up @@ -677,20 +680,9 @@ class _ArrayOrScalarCommon(
def __rmod__(self, other): ...
def __divmod__(self, other): ...
def __rdivmod__(self, other): ...
def __lshift__(self, other): ...
def __rlshift__(self, other): ...
def __rshift__(self, other): ...
def __rrshift__(self, other): ...
def __and__(self, other): ...
def __rand__(self, other): ...
def __xor__(self, other): ...
def __rxor__(self, other): ...
def __or__(self, other): ...
def __ror__(self, other): ...
def __neg__(self: _ArraySelf) -> _ArraySelf: ...
def __pos__(self: _ArraySelf) -> _ArraySelf: ...
def __abs__(self: _ArraySelf) -> _ArraySelf: ...
def __invert__(self: _ArraySelf) -> _ArraySelf: ...
def astype(
self: _ArraySelf,
dtype: DtypeLike,
Expand Down Expand Up @@ -1257,6 +1249,17 @@ class ndarray(_ArrayOrScalarCommon, Iterable, Sized, Container):
def __rpow__(self, other: ArrayLike) -> Union[ndarray, generic]: ...
def __truediv__(self, other: ArrayLike) -> Union[ndarray, generic]: ...
def __rtruediv__(self, other: ArrayLike) -> Union[ndarray, generic]: ...
def __invert__(self: _ArraySelf) -> Union[_ArraySelf, integer, bool_]: ...
def __lshift__(self, other: ArrayLike) -> Union[ndarray, integer]: ...
def __rlshift__(self, other: ArrayLike) -> Union[ndarray, integer]: ...
def __rshift__(self, other: ArrayLike) -> Union[ndarray, integer]: ...
def __rrshift__(self, other: ArrayLike) -> Union[ndarray, integer]: ...
def __and__(self, other: ArrayLike) -> Union[ndarray, integer, bool_]: ...
def __rand__(self, other: ArrayLike) -> Union[ndarray, integer, bool_]: ...
def __xor__(self, other: ArrayLike) -> Union[ndarray, integer, bool_]: ...
def __rxor__(self, other: ArrayLike) -> Union[ndarray, integer, bool_]: ...
def __or__(self, other: ArrayLike) -> Union[ndarray, integer, bool_]: ...
def __ror__(self, other: ArrayLike) -> Union[ndarray, integer, bool_]: ...
# `np.generic` does not support inplace operations
def __iadd__(self: _ArraySelf, other: ArrayLike) -> _ArraySelf: ...
def __isub__(self: _ArraySelf, other: ArrayLike) -> _ArraySelf: ...
Expand All @@ -1265,11 +1268,11 @@ class ndarray(_ArrayOrScalarCommon, Iterable, Sized, Container):
def __ifloordiv__(self: _ArraySelf, other: ArrayLike) -> _ArraySelf: ...
def __ipow__(self: _ArraySelf, other: ArrayLike) -> _ArraySelf: ...
def __imod__(self, other): ...
def __ilshift__(self, other): ...
def __irshift__(self, other): ...
def __iand__(self, other): ...
def __ixor__(self, other): ...
def __ior__(self, other): ...
def __ilshift__(self: _ArraySelf, other: ArrayLike) -> _ArraySelf: ...
def __irshift__(self: _ArraySelf, other: ArrayLike) -> _ArraySelf: ...
def __iand__(self: _ArraySelf, other: ArrayLike) -> _ArraySelf: ...
def __ixor__(self: _ArraySelf, other: ArrayLike) -> _ArraySelf: ...
def __ior__(self: _ArraySelf, other: ArrayLike) -> _ArraySelf: ...

# NOTE: while `np.generic` is not technically an instance of `ABCMeta`,
# the `@abstractmethod` decorator is herein used to (forcefully) deny
Expand Down Expand Up @@ -1329,6 +1332,17 @@ class bool_(generic):
__rpow__: _BoolOp[int8]
__truediv__: _BoolTrueDiv
__rtruediv__: _BoolTrueDiv
def __invert__(self) -> bool_: ...
__lshift__: _BoolBitOp[int8]
__rlshift__: _BoolBitOp[int8]
__rshift__: _BoolBitOp[int8]
__rrshift__: _BoolBitOp[int8]
__and__: _BoolBitOp[bool_]
__rand__: _BoolBitOp[bool_]
__xor__: _BoolBitOp[bool_]
__rxor__: _BoolBitOp[bool_]
__or__: _BoolBitOp[bool_]
__ror__: _BoolBitOp[bool_]

class object_(generic):
def __init__(self, __value: object = ...) -> None: ...
Expand Down Expand Up @@ -1374,6 +1388,18 @@ class integer(number): # type: ignore
def __index__(self) -> int: ...
__truediv__: _IntTrueDiv
__rtruediv__: _IntTrueDiv
def __invert__(self: _IntType) -> _IntType: ...
# Ensure that objects annotated as `integer` support bit-wise operations
def __lshift__(self, other: Union[_IntLike, _BoolLike]) -> integer: ...
def __rlshift__(self, other: Union[_IntLike, _BoolLike]) -> integer: ...
def __rshift__(self, other: Union[_IntLike, _BoolLike]) -> integer: ...
def __rrshift__(self, other: Union[_IntLike, _BoolLike]) -> integer: ...
def __and__(self, other: Union[_IntLike, _BoolLike]) -> integer: ...
def __rand__(self, other: Union[_IntLike, _BoolLike]) -> integer: ...
def __or__(self, other: Union[_IntLike, _BoolLike]) -> integer: ...
def __ror__(self, other: Union[_IntLike, _BoolLike]) -> integer: ...
def __xor__(self, other: Union[_IntLike, _BoolLike]) -> integer: ...
def __rxor__(self, other: Union[_IntLike, _BoolLike]) -> integer: ...

class signedinteger(integer): # type: ignore
__add__: _SignedIntOp
Expand All @@ -1386,6 +1412,16 @@ class signedinteger(integer): # type: ignore
__rfloordiv__: _SignedIntOp
__pow__: _SignedIntOp
__rpow__: _SignedIntOp
__lshift__: _SignedIntBitOp
__rlshift__: _SignedIntBitOp
__rshift__: _SignedIntBitOp
__rrshift__: _SignedIntBitOp
__and__: _SignedIntBitOp
__rand__: _SignedIntBitOp
__xor__: _SignedIntBitOp
__rxor__: _SignedIntBitOp
__or__: _SignedIntBitOp
__ror__: _SignedIntBitOp

class int8(signedinteger):
def __init__(self, __value: _IntValue = ...) -> None: ...
Expand Down Expand Up @@ -1429,6 +1465,16 @@ class unsignedinteger(integer): # type: ignore
__rfloordiv__: _UnsignedIntOp
__pow__: _UnsignedIntOp
__rpow__: _UnsignedIntOp
__lshift__: _UnsignedIntBitOp
__rlshift__: _UnsignedIntBitOp
__rshift__: _UnsignedIntBitOp
__rrshift__: _UnsignedIntBitOp
__and__: _UnsignedIntBitOp
__rand__: _UnsignedIntBitOp
__xor__: _UnsignedIntBitOp
__rxor__: _UnsignedIntBitOp
__or__: _UnsignedIntBitOp
__ror__: _UnsignedIntBitOp

class uint8(unsignedinteger):
def __init__(self, __value: _IntValue = ...) -> None: ...
Expand Down Expand Up @@ -1458,6 +1504,7 @@ class floating(inexact): # type: ignore
__pow__: _FloatOp
__rpow__: _FloatOp

_IntType = TypeVar("_IntType", bound=integer)
_FloatType = TypeVar('_FloatType', bound=floating)

class float16(floating):
Expand Down
24 changes: 24 additions & 0 deletions numpy/typing/_callable.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@
HAVE_PROTOCOL = True

if HAVE_PROTOCOL:
_IntType = TypeVar("_IntType", bound=integer)
_NumberType = TypeVar("_NumberType", bound=number)
_NumberType_co = TypeVar("_NumberType_co", covariant=True, bound=number)
_GenericType_co = TypeVar("_GenericType_co", covariant=True, bound=generic)
Expand All @@ -61,6 +62,14 @@ def __call__(self, __other: complex) -> complex128: ...
@overload
def __call__(self, __other: _NumberType) -> _NumberType: ...

class _BoolBitOp(Protocol[_GenericType_co]):
@overload
def __call__(self, __other: _BoolLike) -> _GenericType_co: ...
@overload # platform dependent
def __call__(self, __other: int) -> Union[int32, int64]: ...
@overload
def __call__(self, __other: _IntType) -> _IntType: ...

class _BoolSub(Protocol):
# Note that `__other: bool_` is absent here
@overload # platform dependent
Expand Down Expand Up @@ -103,6 +112,15 @@ def __call__(self, __other: float) -> floating: ...
@overload
def __call__(self, __other: complex) -> complexfloating[floating]: ...

class _UnsignedIntBitOp(Protocol):
# TODO: The likes of `uint64 | np.signedinteger` will fail as there
# is no signed integer type large enough to hold a `uint64`
# See https://github.com/numpy/numpy/issues/2524
Comment on lines +116 to +118
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comment seems misplaced to me, it applies to the second overload only.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is correct. Inserting a comment between two overloads looks visually rather bad though, hence why it is on top.

@overload
def __call__(self, __other: Union[bool, unsignedinteger]) -> unsignedinteger: ...
@overload
def __call__(self, __other: Union[int, signedinteger]) -> signedinteger: ...

class _SignedIntOp(Protocol):
@overload
def __call__(self, __other: Union[int, signedinteger]) -> signedinteger: ...
Expand All @@ -111,6 +129,9 @@ def __call__(self, __other: float) -> floating: ...
@overload
def __call__(self, __other: complex) -> complexfloating[floating]: ...

class _SignedIntBitOp(Protocol):
def __call__(self, __other: Union[int, signedinteger]) -> signedinteger: ...

class _FloatOp(Protocol):
@overload
def __call__(self, __other: _FloatLike) -> floating: ...
Expand All @@ -125,12 +146,15 @@ def __call__(self, __other: _NumberLike) -> number: ...

else:
_BoolOp = Any
_BoolBitOp = Any
_BoolSub = Any
_BoolTrueDiv = Any
_TD64Div = Any
_IntTrueDiv = Any
_UnsignedIntOp = Any
_UnsignedIntBitOp = Any
_SignedIntOp = Any
_SignedIntBitOp = Any
_FloatOp = Any
_ComplexOp = Any
_NumberOp = Any
20 changes: 20 additions & 0 deletions numpy/typing/tests/data/fail/bitwise_ops.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
import numpy as np

i8 = np.int64()
i4 = np.int32()
u8 = np.uint64()
b_ = np.bool_()
i = int()

f8 = np.float64()

b_ >> f8 # E: No overload variant
i8 << f8 # E: incompatible type
i | f8 # E: Unsupported operand types
i8 ^ f8 # E: incompatible type
u8 & f8 # E: No overload variant
~f8 # E: Unsupported operand type

# mypys' error message for `NoReturn` is unfortunately pretty bad
# TODO: Reenable this once we add support for numerical precision for `number`s
# a = u8 | 0 # E: Need type annotation
131 changes: 131 additions & 0 deletions numpy/typing/tests/data/pass/bitwise_ops.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,131 @@
import numpy as np

i8 = np.int64(1)
u8 = np.uint64(1)

< F438 /td> i4 = np.int32(1)
u4 = np.uint32(1)

b_ = np.bool_(1)

b = bool(1)
i = int(1)

AR = np.array([0, 1, 2], dtype=np.int32)
AR.setflags(write=False)


i8 << i8
i8 >> i8
i8 | i8
i8 ^ i8
i8 & i8

i8 << AR
i8 >> AR
i8 | AR
i8 ^ AR
i8 & AR

i4 << i4
i4 >> i4
i4 | i4
i4 ^ i4
i4 & i4

i8 << i4
i8 >> i4
i8 | i4
i8 ^ i4
i8 & i4

i8 << i
i8 >> i
i8 | i
i8 ^ i
i8 & i

i8 << b_
i8 >> b_
i8 | b_
i8 ^ b_
i8 & b_

i8 << b
i8 >> b
i8 | b
i8 ^ b
i8 & b

u8 << u8
u8 >> u8
u8 | u8
u8 ^ u8
u8 & u8

u8 << AR
u8 >> AR
u8 | AR
u8 ^ AR
u8 & AR

u4 << u4
u4 >> u4
u4 | u4
u4 ^ u4
u4 & u4

u4 << i4
u4 >> i4
u4 | i4
u4 ^ i4
u4 & i4

u4 << i
u4 >> i
u4 | i
u4 ^ i
u4 & i

u8 << b_
u8 >> b_
u8 | b_
u8 ^ b_
u8 & b_

u8 << b
u8 >> b
u8 | b
u8 ^ b
u8 & b

b_ << b_
b_ >> b_
b_ | b_
b_ ^ b_
b_ & b_

b_ << AR
b_ >> AR
b_ | AR
b_ ^ AR
b_ & AR

b_ << b
b_ >> b
b_ | b
b_ ^ b
b_ & b

b_ << i
b_ >> i
b_ | i
b_ ^ i
b_ & i

~i8
~i4
~u8
~u4
~b_
~AR
Loading
0