8000 Merge pull request #27420 from jorenham/typing/complexfloating-option… · numpy/numpy@88a649b · GitHub
[go: up one dir, main page]
Skip to content

Commit 88a649b

Browse files
charrischarris
authored
Merge pull request #27420 from jorenham/typing/complexfloating-optional-nbit2
TYP: Optional 2nd ``numpy.complexfloating`` type parameter
2 parents e62b3b7 + 1763445 commit 88a649b

File tree

3 files changed

+73
-39
lines changed

3 files changed

+73
-39
lines changed

‎doc/release/upcoming_changes/27420.new_feature.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
* ``np.complexfloating[T, T]`` can now also be written as
2+
``np.complexfloating[T]``

‎numpy/__init__.pyi

Lines changed: 71 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,6 @@ from numpy._typing._callable import (
161161
_FloatOp,
162162
_FloatMod,
163163
_FloatDivMod,
164-
_ComplexOp,
165164
_NumberOp,
166165
_ComparisonOpLT,
167166
_ComparisonOpLE,
@@ -199,13 +198,10 @@ from typing import (
199198
Literal as L,
200199
Any,
201200
Generator,
202-
Generic,
203201
NoReturn,
204-
overload,
205202
SupportsComplex,
206203
SupportsFloat,
207204
SupportsInt,
208-
Protocol,
209205
SupportsIndex,
210206
Final,
211207
final,
@@ -218,7 +214,7 @@ from typing import (
218214
# This is because the `typeshed` stubs for the standard library include
219215
# `typing_extensions` stubs:
220216
# https://github.com/python/typeshed/blob/main/stdlib/typing_extensions.pyi
221-
from typing_extensions import LiteralString, Self, TypeVar
217+
from typing_extensions import Generic, LiteralString, Protocol, Self, TypeVar, overload
222218

223219
from numpy import (
224220
core,
@@ -3086,8 +3082,9 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
30863082
# See https://github.com/numpy/numpy-stubs/pull/80 for more details.
30873083

30883084
_ScalarType = TypeVar("_ScalarType", bound=generic)
3085+
_NBit = TypeVar("_NBit", bound=NBitBase)
30893086
_NBit1 = TypeVar("_NBit1", bound=NBitBase)
3090-
_NBit2 = TypeVar("_NBit2", bound=NBitBase)
3087+
_NBit2 = TypeVar("_NBit2", bound=NBitBase, default=_NBit1)
30913088

30923089
class generic(_ArrayOrScalarCommon):
30933090
@abstractmethod
@@ -3585,7 +3582,6 @@ float32: TypeAlias = floating[_32Bit]
35853582

35863583
# NOTE: `_64Bit` is equivalent to `_64Bit | _32Bit | _16Bit | _8Bit`
35873584
_Float64_co: TypeAlias = float | floating[_64Bit] | integer[_64Bit] | np.bool
3588-
_Complex128_co: TypeAlias = complex | complexfloating[_64Bit, _64Bit] | _Float64_co
35893585

35903586
# either a C `double`, `float`, or `longdouble`
35913587
class float64(floating[_64Bit], float): # type: ignore[misc]
@@ -3714,6 +3710,9 @@ single: TypeAlias = floating[_NBitSingle]
37143710
double: TypeAlias = floating[_NBitDouble]
37153711
longdouble: TypeAlias = floating[_NBitLongDouble]
37163712

3713+
_Complex64_co: TypeAlias = builtins.bool | np.bool | number[_32Bit]
3714+
_Complex128_co: TypeAlias = complex | np.bool | number[_64Bit]
3715+
37173716
# The main reason for `complexfloating` having two typevars is cosmetic.
37183717
# It is used to clarify why `complex128`s precision is `_64Bit`, the latter
37193718
# describing the two 64 bit floats representing its real and imaginary component
@@ -3726,19 +3725,73 @@ class complexfloating(inexact[_NBit1], Generic[_NBit1, _NBit2]):
37263725
def real(self) -> floating[_NBit1]: ... # type: ignore[override]
37273726
@property
37283727
def imag(self) -> floating[_NBit2]: ... # type: ignore[override]
3729-
def __abs__(self) -> floating[_NBit1]: ... # type: ignore[override]
3728+
def __abs__(self) -> floating[_NBit1 | _NBit2]: ... # type: ignore[override]
37303729
# NOTE: Deprecated
37313730
# def __round__(self, ndigits=...): ...
3732-
__add__: _ComplexOp[_NBit1]
3733-
__radd__: _ComplexOp[_NBit1]
3734-
__sub__: _ComplexOp[_NBit1]
3735-
__rsub__: _ComplexOp[_NBit1]
3736-
6DB6 __mul__: _ComplexOp[_NBit1]
3737-
__rmul__: _ComplexOp[_NBit1]
3738-
__truediv__: _ComplexOp[_NBit1]
3739-
__rtruediv__: _ComplexOp[_NBit1]
3740-
__pow__: _ComplexOp[_NBit1]
3741-
__rpow__: _ComplexOp[_NBit1]
3731+
@overload
3732+
def __add__(self, other: _Complex64_co, /) -> complexfloating[_NBit1, _NBit2]: ...
3733+
@overload
3734+
def __add__(self, other: complex | float64 | complex128, /) -> complexfloating[_NBit1, _NBit2] | complex128: ...
3735+
@overload
3736+
def __add__(self, other: number[_NBit], /) -> complexfloating[_NBit1, _NBit2] | complexfloating[_NBit, _NBit]: ...
3737+
@overload
3738+
def __radd__(self, other: _Complex64_co, /) -> complexfloating[_NBit1, _NBit2]: ...
3739+
@overload
3740+
def __radd__(self, other: complex, /) -> complexfloating[_NBit1, _NBit2] | complex128: ...
3741+
@overload
3742+
def __radd__(self, other: number[_NBit], /) -> complexfloating[_NBit1, _NBit2] | complexfloating[_NBit, _NBit]: ...
3743+
3744+
@overload
3745+
def __sub__(self, other: _Complex64_co, /) -> complexfloating[_NBit1, _NBit2]: ...
3746+
@overload
3747+
def __sub__(self, other: complex | float64 | complex128, /) -> complexfloating[_NBit1, _NBit2] | complex128: ...
3748+
@overload
3749+
def __sub__(self, other: number[_NBit], /) -> complexfloating[_NBit1, _NBit2] | complexfloating[_NBit, _NBit]: ...
3750+
@overload
3751+
def __rsub__(self, other: _Complex64_co, /) -> complexfloating[_NBit1, _NBit2]: ...
3752+
@overload
3753+
def __rsub__(self, other: complex, /) -> complexfloating[_NBit1, _NBit2] | complex128: ...
3754+
@overload
3755+
def __rsub__(self, other: number[_NBit], /) -> complexfloating[_NBit1, _NBit2] | complexfloating[_NBit, _NBit]: ...
3756+
3757+
@overload
3758+
def __mul__(self, other: _Complex64_co, /) -> complexfloating[_NBit1, _NBit2]: ...
3759+
@overload
3760+
def __mul__(self, other: complex | float64 | complex128, /) -> complexfloating[_NBit1, _NBit2] | complex128: ...
3761+
@overload
3762+
def __mul__(self, other: number[_NBit], /) -> complexfloating[_NBit1, _NBit2] | complexfloating[_NBit, _NBit]: ...
3763+
@overload
3764+
def __rmul__(self, other: _Complex64_co, /) -> complexfloating[_NBit1, _NBit2]: ...
3765+
@overload
3766+
def __rmul__(self, other: complex, /) -> complexfloating[_NBit1, _NBit2] | complex128: ...
3767+
@overload
3768+
def __rmul__(self, other: number[_NBit], /) -> complexfloating[_NBit1, _NBit2] | complexfloating[_NBit, _NBit]: ...
3769+
3770+
@overload
3771+
def __truediv__(self, other: _Complex64_co, /) -> complexfloating[_NBit1, _NBit2]: ...
3772+
@overload
3773+
def __truediv__(self, other: complex | float64 | complex128, /) -> complexfloating[_NBit1, _NBit2] | complex128: ...
3774+
@overload
3775+
def __truediv__(self, other: number[_NBit], /) -> complexfloating[_NBit1, _NBit2] | complexfloating[_NBit, _NBit]: ...
3776+
@overload
3777+
def __rtruediv__(self, other: _Complex64_co, /) -> complexfloating[_NBit1, _NBit2]: ...
3778+
@overload
3779+
def __rtruediv__(self, other: complex, /) -> complexfloating[_NBit1, _NBit2] | complex128: ...
3780+
@overload
3781+
def __rtruediv__(self, other: number[_NBit], /) -> complexfloating[_NBit1, _NBit2] | complexfloating[_NBit, _NBit]: ...
3782+
3783+
@overload
3784+
def __pow__(self, other: _Complex64_co, /) -> complexfloating[_NBit1, _NBit2]: ...
3785+
@overload
3786+
def __pow__(self, other: complex | float64 | complex128, /) -> complexfloating[_NBit1, _NBit2] | complex128: ...
3787+
@overload
3788+
def __pow__(self, other: number[_NBit], /) -> complexfloating[_NBit1, _NBit2] | complexfloating[_NBit, _NBit]: ...
3789+
@overload
3790+
def __rpow__(self, other: _Complex64_co, /) -> complexfloating[_NBit1, _NBit2]: ...
3791+
@overload
3792+
def __rpow__(self, other: complex, /) -> complexfloating[_NBit1, _NBit2] | complex128: ...
3793+
@overload
3794+
def __rpow__(self, other: number[_NBit], /) -> complexfloating[_NBit1, _NBit2] | complexfloating[_NBit, _NBit]: ...
37423795

37433796
complex64: TypeAlias = complexfloating[_32Bit, _32Bit]
37443797

‎numpy/_typing/_callable.pyi

Lines changed: 0 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -298,27 +298,6 @@ class _FloatDivMod(Protocol[_NBit1]):
298298
self, other: integer[_NBit2] | floating[_NBit2], /
299299
) -> _2Tuple[floating[_NBit1]] | _2Tuple[floating[_NBit2]]: ...
300300

301-
class _ComplexOp(Protocol[_NBit1]):
302-
@overload
303-
def __call__(self, other: bool, /) -> complexfloating[_NBit1, _NBit1]: ...
304-
@overload
305-
def __call__(
306-
self, other: int, /
307-
) -> complexfloating[_NBit1, _NBit1] | complexfloating[_NBitInt, _NBitInt]: ...
308-
@overload
309-
def __call__(
310-
self, other: complex, /,
311-
) -> complexfloating[_NBit1, _NBit1] | complex128: ...
312-
@overload
313-
def __call__(
314-
self,
315-
other: (
316-
integer[_NBit2]
317-
| floating[_NBit2]
318-
| complexfloating[_NBit2, _NBit2]
319-
), /,
320-
) -> complexfloating[_NBit1, _NBit1] | complexfloating[_NBit2, _NBit2]: ...
321-
322301
class _NumberOp(Protocol):
323302
def __call__(self, other: _NumberLike_co, /) -> Any: ...
324303

0 commit comments

Comments
 (0)
0