Skip to content

Commit 0c899da

Browse files
committed
♻️ move the nep50 builtin scalar protocols from numpy,__init__ to _numtype
1 parent fff5c9c commit 0c899da

File tree

3 files changed

+42
-41
lines changed

3 files changed

+42
-41
lines changed

src/_numtype/__init__.pyi

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,8 @@ from ._nep50 import (
5151
CanCast0D as CanCast0D,
5252
CanCastND as CanCastND,
5353
CanNEP50 as CanNEP50,
54+
CanNEP50Complex as CanNEP50Complex,
55+
CanNEP50Float as CanNEP50Float,
5456
MatchND as MatchND,
5557
PromoteWith as PromoteWith,
5658
PromoteWith0D as PromoteWith0D,

src/_numtype/_nep50.pyi

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ from typing_extensions import TypeAliasType, TypeVar
77

88
import numpy as np
99

10-
__all__ = ["CanCast0D", "CanCastND", "CanNEP50", "MatchND", "PromoteWith", "PromoteWith0D"]
10+
__all__ = ["CanCast0D", "CanCastND", "CanNEP50", "CanNEP50Complex", "CanNEP50Float", "MatchND", "PromoteWith", "PromoteWith0D"]
1111

1212
_T_co = TypeVar("_T_co", covariant=True)
1313
_BelowT_contra = TypeVar("_BelowT_contra", bound=np.generic, contravariant=True)
@@ -16,6 +16,9 @@ _OtherT_contra = TypeVar("_OtherT_contra", contravariant=True)
1616
_MatchT_co = TypeVar("_MatchT_co", bound=np.generic, covariant=True, default=Any)
1717
_ShapeT_co = TypeVar("_ShapeT_co", bound=tuple[int, ...], covariant=True)
1818

19+
_ComplexFloatingT_co = TypeVar("_ComplexFloatingT_co", bound=np.complexfloating, covariant=True)
20+
_InexactT_co = TypeVar("_InexactT_co", bound=np.inexact, covariant=True)
21+
1922
@type_check_only
2023
class CanNEP50(Protocol[_BelowT_contra, _AboveT_contra, _MatchT_co]):
2124
def __nep50__(self, below: _BelowT_contra, above: _AboveT_contra, /) -> _MatchT_co: ...
@@ -48,6 +51,14 @@ class _CanNEP50Rule5(Protocol[_OtherT_contra, _T_co]):
4851
class _CanNEP50Rule6(Protocol[_OtherT_contra, _T_co]):
4952
def __nep50_rule6__(self, other: _OtherT_contra, /) -> _T_co: ...
5053

54+
@type_check_only
55+
class CanNEP50Float(Protocol[_InexactT_co]):
56+
def __nep50_float__(self, /) -> _InexactT_co: ...
57+
58+
@type_check_only
59+
class CanNEP50Complex(Protocol[_ComplexFloatingT_co]):
60+
def __nep50_complex__(self, /) -> _ComplexFloatingT_co: ...
61+
5162
_WithT = TypeVar("_WithT")
5263
_OutT = TypeVar("_OutT", bound=np.generic)
5364

src/numpy-stubs/__init__.pyi

Lines changed: 28 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -628,9 +628,7 @@ _RealScalarT = TypeVar("_RealScalarT", bound=bool_ | integer | floating | object
628628
_IntegerT = TypeVar("_IntegerT", bound=integer)
629629
_FloatingT = TypeVar("_FloatingT", bound=floating)
630630
_ComplexFloatingT = TypeVar("_ComplexFloatingT", bound=complexfloating)
631-
_ComplexFloatingT_co = TypeVar("_ComplexFloatingT_co", bound=complexfloating, covariant=True)
632631
_InexactT = TypeVar("_InexactT", bound=inexact)
633-
_InexactT_co = TypeVar("_InexactT_co", bound=inexact, covariant=True)
634632
_NumberT = TypeVar("_NumberT", bound=number)
635633
_CharT = TypeVar("_CharT", bound=character)
636634

@@ -653,6 +651,8 @@ _TD64ItemT_co = TypeVar("_TD64ItemT_co", bound=dt.timedelta | int | None, defaul
653651
_DT64ItemT_co = TypeVar("_DT64ItemT_co", bound=dt.date | int | None, default=dt.date | int | None, covariant=True)
654652
_TD64UnitT = TypeVar("_TD64UnitT", bound=_TD64Unit, default=_TD64Unit)
655653

654+
_FloatSizeT_co = TypeVar("_FloatSizeT_co", bound=L[2, 4, 8, 12, 16], covariant=True)
655+
656656
###
657657
# Type Aliases (for internal use only)
658658

@@ -4328,14 +4328,6 @@ bool = bool_
43284328

43294329
# TODO(jorenham): Move these protocols to _numtype
43304330

4331-
@type_check_only
4332-
class _CanNEP50Float(Protocol[_InexactT_co]):
4333-
def __nep50_float__(self, /) -> _InexactT_co: ...
4334-
4335-
@type_check_only
4336-
class _CanNEP50Complex(Protocol[_ComplexFloatingT_co]):
4337-
def __nep50_complex__(self, /) -> _ComplexFloatingT_co: ...
4338-
43394331
class number(
43404332
_CmpOpMixin[_nt.CoComplex_0d, _nt.CoComplex_1nd],
43414333
generic[_NumberItemT_co],
@@ -4368,95 +4360,95 @@ class number(
43684360
@overload
43694361
def __add__(self, x: _nt.CanCast0D[Self] | int, /) -> Self: ...
43704362
@overload
4371-
def __add__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
4363+
def __add__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
43724364
@overload
4373-
def __add__(self: _CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
4365+
def __add__(self: _nt.CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
43744366
@overload
43754367
def __add__(self, x: _nt.PromoteWith0D[Self, _ScalarT], /) -> _ScalarT: ...
43764368

43774369
# keep in sync with __add__
43784370
@overload
43794371
def __radd__(self, x: _nt.CanCast0D[Self] | int, /) -> Self: ...
43804372
@overload
4381-
def __radd__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
4373+
def __radd__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
43824374
@overload
4383-
def __radd__(self: _CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
4375+
def __radd__(self: _nt.CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
43844376
@overload
43854377
def __radd__(self, x: _nt.PromoteWith0D[Self, _ScalarT], /) -> _ScalarT: ...
43864378

43874379
# keep in sync with __add__
43884380
@overload
43894381
def __sub__(self, x: _nt.CanCast0D[Self] | int, /) -> Self: ...
43904382
@overload
4391-
def __sub__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
4383+
def __sub__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
43924384
@overload
4393-
def __sub__(self: _CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
4385+
def __sub__(self: _nt.CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
43944386
@overload
43954387
def __sub__(self, x: _nt.PromoteWith0D[Self, _ScalarT], /) -> _ScalarT: ...
43964388

43974389
# keep in sync with __add__
43984390
@overload
43994391
def __rsub__(self, x: _nt.CanCast0D[Self] | int, /) -> Self: ...
44004392
@overload
4401-
def __rsub__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
4393+
def __rsub__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
44024394
@overload
4403-
def __rsub__(self: _CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
4395+
def __rsub__(self: _nt.CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
44044396
@overload
44054397
def __rsub__(self, x: _nt.PromoteWith0D[Self, _ScalarT], /) -> _ScalarT: ...
44064398

44074399
# keep in sync with __add__
44084400
@overload
44094401
def __mul__(self, x: _nt.CanCast0D[Self] | int, /) -> Self: ...
44104402
@overload
4411-
def __mul__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
4403+
def __mul__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
44124404
@overload
4413-
def __mul__(self: _CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
4405+
def __mul__(self: _nt.CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
44144406
@overload
44154407
def __mul__(self, x: _nt.PromoteWith0D[Self, _ScalarT], /) -> _ScalarT: ...
44164408

44174409
# keep in sync with __add__
44184410
@overload
44194411
def __rmul__(self, x: _nt.CanCast0D[Self] | int, /) -> Self: ...
44204412
@overload
4421-
def __rmul__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
4413+
def __rmul__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
44224414
@overload
4423-
def __rmul__(self: _CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
4415+
def __rmul__(self: _nt.CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
44244416
@overload
44254417
def __rmul__(self, x: _nt.PromoteWith0D[Self, _ScalarT], /) -> _ScalarT: ...
44264418

44274419
# keep in sync with __add__
44284420
@overload
44294421
def __pow__(self, x: _nt.CanCast0D[Self] | int, mod: None = None, /) -> Self: ...
44304422
@overload
4431-
def __pow__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, mod: None = None, /) -> _InexactT: ...
4423+
def __pow__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, mod: None = None, /) -> _InexactT: ...
44324424
@overload
4433-
def __pow__(self: _CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, mod: None = None, /) -> _ComplexFloatingT: ...
4425+
def __pow__(self: _nt.CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, mod: None = None, /) -> _ComplexFloatingT: ...
44344426
@overload
44354427
def __pow__(self, x: _nt.PromoteWith0D[Self, _ScalarT], mod: None = None, /) -> _ScalarT: ...
44364428

44374429
# keep in sync with __add__
44384430
@overload
44394431
def __rpow__(self, x: _nt.CanCast0D[Self] | int, mod: None = None, /) -> Self: ...
44404432
@overload
4441-
def __rpow__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, mod: None = None, /) -> _InexactT: ...
4433+
def __rpow__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, mod: None = None, /) -> _InexactT: ...
44424434
@overload
4443-
def __rpow__(self: _CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, mod: None = None, /) -> _ComplexFloatingT: ...
4435+
def __rpow__(self: _nt.CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, mod: None = None, /) -> _ComplexFloatingT: ...
44444436
@overload
44454437
def __rpow__(self, x: _nt.PromoteWith0D[Self, _ScalarT], mod: None = None, /) -> _ScalarT: ...
44464438

44474439
#
44484440
@overload
44494441
def __truediv__(self, x: _nt.CoFloating_0d | _nt.CanCast0D[Self] | _nt.Just[number], /) -> inexact: ...
44504442
@overload
4451-
def __truediv__(self: _CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
4443+
def __truediv__(self: _nt.CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
44524444
@overload
44534445
def __truediv__(self, x: _nt.PromoteWith0D[Self, _InexactT], /) -> _InexactT: ...
44544446

44554447
#
44564448
@overload
44574449
def __rtruediv__(self, x: _nt.CoFloating_0d | _nt.CanCast0D[Self] | _nt.Just[number], /) -> inexact: ...
44584450
@overload
4459-
def __rtruediv__(self: _CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
4451+
def __rtruediv__(self: _nt.CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
44604452
@overload
44614453
def __rtruediv__(self, x: _nt.PromoteWith0D[Self, _InexactT], /) -> _InexactT: ...
44624454

@@ -4516,44 +4508,44 @@ class integer(_IntegralMixin, _RoundMixin, number[_BitT, int]):
45164508
@overload
45174509
def __floordiv__(self, x: _nt.CanCast0D[Self] | int, /) -> Self: ...
45184510
@overload
4519-
def __floordiv__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
4511+
def __floordiv__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
45204512
@overload
45214513
def __floordiv__(self, x: _nt.PromoteWith0D[Self, _RealScalarT], /) -> _RealScalarT: ...
45224514
#
45234515
@overload
45244516
def __rfloordiv__(self, x: _nt.CanCast0D[Self] | int, /) -> Self: ...
45254517
@overload
4526-
def __rfloordiv__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
4518+
def __rfloordiv__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
45274519
@overload
45284520
def __rfloordiv__(self, x: _nt.PromoteWith0D[Self, _RealScalarT], /) -> _RealScalarT: ...
45294521

45304522
#
45314523
@overload
45324524
def __mod__(self, x: _nt.CanCast0D[Self] | int, /) -> Self: ...
45334525
@overload
4534-
def __mod__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
4526+
def __mod__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
45354527
@overload
45364528
def __mod__(self, x: _nt.PromoteWith0D[Self, _RealScalarT], /) -> _RealScalarT: ...
45374529
#
45384530
@overload
45394531
def __rmod__(self, x: _nt.CanCast0D[Self] | int, /) -> Self: ...
45404532
@overload
4541-
def __rmod__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
4533+
def __rmod__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _InexactT: ...
45424534
@overload
45434535
def __rmod__(self, x: _nt.PromoteWith0D[Self, _RealScalarT], /) -> _RealScalarT: ...
45444536

45454537
#
45464538
@overload
45474539
def __divmod__(self, x: _nt.CanCast0D[Self] | int, /) -> _2Tuple[Self]: ...
45484540
@overload
4549-
def __divmod__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _2Tuple[_InexactT]: ...
4541+
def __divmod__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _2Tuple[_InexactT]: ...
45504542
@overload
45514543
def __divmod__(self, x: _nt.PromoteWith0D[Self, _RealScalarT], /) -> _2Tuple[_RealScalarT]: ...
45524544
#
45534545
@overload
45544546
def __rdivmod__(self, x: _nt.CanCast0D[Self] | int, /) -> _2Tuple[Self]: ...
45554547
@overload
4556-
def __rdivmod__(self: _CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _2Tuple[_InexactT]: ...
4548+
def __rdivmod__(self: _nt.CanNEP50Float[_InexactT], x: _nt.JustFloat, /) -> _2Tuple[_InexactT]: ...
45574549
@overload
45584550
def __rdivmod__(self, x: _nt.PromoteWith0D[Self, _RealScalarT], /) -> _2Tuple[_RealScalarT]: ...
45594551

@@ -4988,14 +4980,14 @@ class inexact(number[_BitT, _InexactItemT_co], Generic[_BitT, _InexactItemT_co])
49884980
@overload
49894981
def __truediv__(self, x: int | _nt.JustFloat | _nt.CanCast0D[Self], /) -> Self: ...
49904982
@overload
4991-
def __truediv__(self: _CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
4983+
def __truediv__(self: _nt.CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
49924984
@overload
49934985
def __truediv__(self, x: _nt.PromoteWith0D[Self, _InexactT], /) -> _InexactT: ... # pyright: ignore[reportIncompatibleMethodOverride]
49944986
#
49954987
@overload
49964988
def __rtruediv__(self, x: int | _nt.JustFloat | _nt.CanCast0D[Self], /) -> Self: ...
49974989
@overload
4998-
def __rtruediv__(self: _CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
4990+
def __rtruediv__(self: _nt.CanNEP50Complex[_ComplexFloatingT], x: _nt.JustComplex, /) -> _ComplexFloatingT: ...
49994991
@overload
50004992
def __rtruediv__(self, x: _nt.PromoteWith0D[Self, _InexactT], /) -> _InexactT: ... # pyright: ignore[reportIncompatibleMethodOverride]
50014993

@@ -5050,16 +5042,12 @@ class floating(_RealMixin, _RoundMixin, inexact[_BitT, float]):
50505042
@overload
50515043
def __rdivmod__(self, x: _nt.PromoteWith0D[Self, _FloatingT], /) -> _2Tuple[_FloatingT]: ...
50525044

5053-
_FloatSizeT_co = TypeVar("_FloatSizeT_co", bound=L[2, 4, 8, 12, 16], covariant=True)
5054-
50555045
@type_check_only
50565046
class _FloatMixin(Generic[_FloatSizeT_co]):
50575047
@property
50585048
def itemsize(self) -> _FloatSizeT_co: ...
50595049
@property
50605050
def nbytes(self) -> _FloatSizeT_co: ...
5061-
5062-
#
50635051
@override
50645052
def __hash__(self, /) -> int: ...
50655053
def is_integer(self, /) -> py_bool: ...

0 commit comments

Comments
 (0)