Skip to content

Further improve return types in the numbers module #11371

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

Closed
wants to merge 6 commits into from
Closed
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
144 changes: 93 additions & 51 deletions stdlib/numbers.pyi
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# Note: these stubs are incomplete. The more complex type
# signatures are currently omitted.
#
# Use SupportsComplex, SupportsFloat and SupportsIndex for return types in this module
# Use _ComplexLike, _RealLike and _IntegralLike for return types in this module
# rather than `numbers.Complex`, `numbers.Real` and `numbers.Integral`,
# to avoid an excessive number of `type: ignore`s in subclasses of these ABCs
# (since type checkers don't see `complex` as a subtype of `numbers.Complex`,
Expand All @@ -10,18 +10,56 @@
import sys
from _typeshed import Incomplete
from abc import ABCMeta, abstractmethod
from typing import Literal, SupportsFloat, SupportsIndex, overload
from typing import Literal, Protocol, overload
from typing_extensions import TypeAlias

__all__ = ["Number", "Complex", "Real", "Rational", "Integral"]

############################
# Protocols for return types
############################

# `_ComplexLike` is a structural-typing approximation
# of the `Complex` ABC, which is not (and cannot be) a protocol
class _SupportsComplex(Protocol):
def __complex__(self) -> complex: ...
def __abs__(self) -> _RealLike: ...
def __neg__(self) -> _ComplexLike: ...
def __pos__(self) -> _ComplexLike: ...

if sys.version_info >= (3, 11):
from typing import SupportsComplex as _SupportsComplex
_ComplexLike: TypeAlias = _SupportsComplex
else:
# builtins.complex didn't have a __complex__ method on older Pythons
import typing
_ComplexLike: TypeAlias = _SupportsComplex | complex

# _RealLike is a structural-typing approximation
# of the `Real` ABC, which is not (and cannot be) a protocol
#
# NOTE: _RealLike can't inherit from _SupportsComplex,
# because not all builtin types that we want to be understood
# as subtypes of _RealLike have a `__complex__` method
# (e.g. `builtins.int.__complex__` does not exist)
class _RealLike(Protocol):
def __neg__(self) -> _ComplexLike: ... # TODO: ideally would be more precise
Copy link
Member

Choose a reason for hiding this comment

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

Why don't you make these two return RealLike?

Copy link
Member Author

Choose a reason for hiding this comment

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

I tried doing this:

--- a/stdlib/numbers.pyi
+++ b/stdlib/numbers.pyi
@@ -41,8 +41,8 @@ else:
 # as subtypes of _RealLike have a `__complex__` method
 # (e.g. `builtins.int.__complex__` does not exist)
 class _RealLike(Protocol):
-    def __neg__(self) -> _ComplexLike: ...  # TODO: ideally would be more precise
-    def __pos__(self) -> _ComplexLike: ...  # TODO: ideally would be more precise
+    def __neg__(self) -> _RealLike: ...
+    def __pos__(self) -> _RealLike: ...
     def __abs__(self) -> _RealLike: ...
     def __float__(self) -> float: ...
     def __trunc__(self) -> _IntegralLike: ...
@@ -143,6 +143,12 @@ class Real(Complex):
     @property
     def imag(self) -> Literal[0]: ...
     def conjugate(self) -> _RealLike: ...  # type: ignore[override]
+    # Not actually overridden at runtime,
+    # but we override these in the stub to give them more precise return types:
+    @abstractmethod
+    def __pos__(self) -> _RealLike: ...
+    @abstractmethod
+    def __neg__(self) -> _RealLike: ...

But on 3.12, mypy complains:

stdlib\numbers.pyi:149: error: Return type "_RealLike" of "__pos__" incompatible with return type "_SupportsComplex" in supertype "Complex"  [override]
stdlib\numbers.pyi:151: error: Return type "_RealLike" of "__neg__" incompatible with return type "_SupportsComplex" in supertype "Complex"  [override]

And on 3.8, mypy complains with:

stdlib\numbers.pyi:149: error: Return type "_RealLike" of "__pos__" incompatible with return type "Union[_SupportsComplex, complex]" in supertype "Complex"  [override]
stdlib\numbers.pyi:151: error: Return type "_RealLike" of "__neg__" incompatible with return type "Union[_SupportsComplex, complex]" in supertype "Complex"  [override]

I think this is because _RealLike does not have a __complex__ method, so isn't seen as a subtype of _SupportsComplex. (It can't have a __complex__ method, because int doesn't have a __complex__ method, and we need int to be understood as a subtype of RealLike.)

This is yet another way in which these ABCs are subtly broken, unfortunately. Another approach might be to return SupportsAbs[_RealLike] | complex (instead of _SupportsComplex | complex) from most of the numbers.Complex methods.

def __pos__(self) -> _ComplexLike: ... # TODO: ideally would be more precise
def __abs__(self) -> _RealLike: ...
def __float__(self) -> float: ...
def __trunc__(self) -> _IntegralLike: ...
def __floor__(self) -> _IntegralLike: ...
def __ceil__(self) -> _IntegralLike: ...

_SupportsComplex: TypeAlias = typing.SupportsComplex | complex
# _IntegralLike is a structural-typing approximation
# of the `Integral` ABC, which is not (and cannot be) a protocol
class _IntegralLike(_RealLike, Protocol):
def __int__(self) -> int: ...
def __index__(self) -> int: ...
def __invert__(self) -> _IntegralLike: ...
def __abs__(self) -> _IntegralLike: ...

__all__ = ["Number", "Complex", "Real", "Rational", "Integral"]
#################
# Module "proper"
#################

class Number(metaclass=ABCMeta):
@abstractmethod
Expand All @@ -35,86 +73,86 @@ class Complex(Number):
def __bool__(self) -> bool: ...
@property
@abstractmethod
def real(self) -> SupportsFloat: ...
def real(self) -> _RealLike: ...
@property
@abstractmethod
def imag(self) -> SupportsFloat: ...
def imag(self) -> _RealLike: ...
@abstractmethod
def __add__(self, other) -> _SupportsComplex: ...
def __add__(self, other) -> _ComplexLike: ...
@abstractmethod
def __radd__(self, other) -> _SupportsComplex: ...
def __radd__(self, other) -> _ComplexLike: ...
@abstractmethod
def __neg__(self) -> _SupportsComplex: ...
def __neg__(self) -> _ComplexLike: ...
@abstractmethod
def __pos__(self) -> _SupportsComplex: ...
def __sub__(self, other) -> _SupportsComplex: ...
def __rsub__(self, other) -> _SupportsComplex: ...
def __pos__(self) -> _ComplexLike: ...
def __sub__(self, other) -> _ComplexLike: ...
def __rsub__(self, other) -> _ComplexLike: ...
@abstractmethod
def __mul__(self, other) -> _SupportsComplex: ...
def __mul__(self, other) -> _ComplexLike: ...
@abstractmethod
def __rmul__(self, other) -> _SupportsComplex: ...
def __rmul__(self, other) -> _ComplexLike: ...
@abstractmethod
def __truediv__(self, other) -> _SupportsComplex: ...
def __truediv__(self, other) -> _ComplexLike: ...
@abstractmethod
def __rtruediv__(self, other) -> _SupportsComplex: ...
def __rtruediv__(self, other) -> _ComplexLike: ...
@abstractmethod
def __pow__(self, exponent) -> _SupportsComplex: ...
def __pow__(self, exponent) -> _ComplexLike: ...
@abstractmethod
def __rpow__(self, base) -> _SupportsComplex: ...
def __rpow__(self, base) -> _ComplexLike: ...
@abstractmethod
def __abs__(self) -> SupportsFloat: ...
def __abs__(self) -> _RealLike: ...
@abstractmethod
def conjugate(self) -> _SupportsComplex: ...
def conjugate(self) -> _ComplexLike: ...
@abstractmethod
def __eq__(self, other: object) -> bool: ...

# See comment at the top of the file
# for why some of these return types are purposefully vague
class Real(Complex, SupportsFloat):
class Real(Complex):
@abstractmethod
def __float__(self) -> float: ...
@abstractmethod
def __trunc__(self) -> SupportsIndex: ...
def __trunc__(self) -> _IntegralLike: ...
@abstractmethod
def __floor__(self) -> SupportsIndex: ...
def __floor__(self) -> _IntegralLike: ...
@abstractmethod
def __ceil__(self) -> SupportsIndex: ...
def __ceil__(self) -> _IntegralLike: ...
@abstractmethod
@overload
def __round__(self, ndigits: None = None) -> SupportsIndex: ...
def __round__(self, ndigits: None = None) -> _IntegralLike: ...
@abstractmethod
@overload
def __round__(self, ndigits: int) -> SupportsFloat: ...
def __divmod__(self, other) -> tuple[SupportsFloat, SupportsFloat]: ...
def __rdivmod__(self, other) -> tuple[SupportsFloat, SupportsFloat]: ...
def __round__(self, ndigits: int) -> _RealLike: ...
def __divmod__(self, other) -> tuple[_RealLike, _RealLike]: ...
def __rdivmod__(self, other) -> tuple[_RealLike, _RealLike]: ...
@abstractmethod
def __floordiv__(self, other) -> SupportsFloat: ...
def __floordiv__(self, other) -> _RealLike: ...
@abstractmethod
def __rfloordiv__(self, other) -> SupportsFloat: ...
def __rfloordiv__(self, other) -> _RealLike: ...
@abstractmethod
def __mod__(self, other) -> SupportsFloat: ...
def __mod__(self, other) -> _RealLike: ...
@abstractmethod
def __rmod__(self, other) -> SupportsFloat: ...
def __rmod__(self, other) -> _RealLike: ...
@abstractmethod
def __lt__(self, other) -> bool: ...
@abstractmethod
def __le__(self, other) -> bool: ...
def __complex__(self) -> complex: ...
@property
def real(self) -> SupportsFloat: ...
def real(self) -> _RealLike: ...
@property
def imag(self) -> Literal[0]: ...
def conjugate(self) -> SupportsFloat: ... # type: ignore[override]
def conjugate(self) -> _RealLike: ... # type: ignore[override]

# See comment at the top of the file
# for why some of these return types are purposefully vague
class Rational(Real):
@property
@abstractmethod
def numerator(self) -> SupportsIndex: ...
def numerator(self) -> _IntegralLike: ...
@property
@abstractmethod
def denominator(self) -> SupportsIndex: ...
def denominator(self) -> _IntegralLike: ...
def __float__(self) -> float: ...

# See comment at the top of the file
Expand All @@ -124,31 +162,35 @@ class Integral(Rational):
def __int__(self) -> int: ...
def __index__(self) -> int: ...
@abstractmethod
def __pow__(self, exponent, modulus: Incomplete | None = None) -> SupportsIndex: ... # type: ignore[override]
def __pow__(self, exponent, modulus: Incomplete | None = None) -> _IntegralLike: ... # type: ignore[override]
@abstractmethod
def __lshift__(self, other) -> SupportsIndex: ...
def __lshift__(self, other) -> _IntegralLike: ...
@abstractmethod
def __rlshift__(self, other) -> SupportsIndex: ...
def __rlshift__(self, other) -> _IntegralLike: ...
@abstractmethod
def __rshift__(self, other) -> SupportsIndex: ...
def __rshift__(self, other) -> _IntegralLike: ...
@abstractmethod
def __rrshift__(self, other) -> SupportsIndex: ...
def __rrshift__(self, other) -> _IntegralLike: ...
@abstractmethod
def __and__(self, other) -> SupportsIndex: ...
def __and__(self, other) -> _IntegralLike: ...
@abstractmethod
def __rand__(self, other) -> SupportsIndex: ...
def __rand__(self, other) -> _IntegralLike: ...
@abstractmethod
def __xor__(self, other) -> SupportsIndex: ...
def __xor__(self, other) -> _IntegralLike: ...
@abstractmethod
def __rxor__(self, other) -> SupportsIndex: ...
def __rxor__(self, other) -> _IntegralLike: ...
@abstractmethod
def __or__(self, other) -> SupportsIndex: ...
def __or__(self, other) -> _IntegralLike: ...
@abstractmethod
def __ror__(self, other) -> SupportsIndex: ...
def __ror__(self, other) -> _IntegralLike: ...
@abstractmethod
def __invert__(self) -> SupportsIndex: ...
def __invert__(self) -> _IntegralLike: ...
def __float__(self) -> float: ...
@property
def numerator(self) -> SupportsIndex: ...
def numerator(self) -> _IntegralLike: ...
@property
def denominator(self) -> Literal[1]: ...
# Not actually overridden at runtime,
# but we override it in the stub to give it a more precise return type:
@abstractmethod
def __abs__(self) -> _IntegralLike: ...