Further improve pow (#6325)

python · Nov 18, 2021 · 48cfe5d · 48cfe5d
1 parent 2f0969c
commit 48cfe5d
Showing 1 changed file with 50 additions and 13 deletions.
diff --git a/stdlib/builtins.pyi b/stdlib/builtins.pyi
@@ -176,6 +176,9 @@ class super(object):
     @overload
     def __init__(self) -> None: ...
 
+_PositiveInteger = Literal[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25]
+_NegativeInteger = Literal[-1, -2, -3, -4, -5, -6, -7, -8, -9, -10, -11, -12, -13, -14, -15, -16, -17, -18, -19, -20]
+
 class int:
     @overload
     def __new__(cls: Type[_T], __x: str | bytes | SupportsInt | SupportsIndex | SupportsTrunc = ...) -> _T: ...
@@ -217,11 +220,15 @@ class int:
     @overload
     def __pow__(self, __x: int, __modulo: Literal[0]) -> NoReturn: ...
     @overload
-    def __pow__(self, __x: Literal[2, 3, 4, 5], __modulo: int | None = ...) -> int: ...
+    def __pow__(self, __x: int, __modulo: int) -> int: ...
+    @overload
+    def __pow__(self, __x: _PositiveInteger, __modulo: None = ...) -> int: ...
+    @overload
+    def __pow__(self, __x: _NegativeInteger, __modulo: None = ...) -> float: ...
     # positive x -> int; negative x -> float
     # return type must be Any as `int | float` causes too many false-positive errors
     @overload
-    def __pow__(self, __x: int, __modulo: int | None = ...) -> Any: ...
+    def __pow__(self, __x: int, __modulo: None = ...) -> Any: ...
     def __rpow__(self, __x: int, __mod: int | None = ...) -> Any: ...
     def __and__(self, __n: int) -> int: ...
     def __or__(self, __n: int) -> int: ...
@@ -1279,44 +1286,74 @@ _M = TypeVar("_M", contravariant=True)
 class _SupportsPow2(Protocol[_E, _T_co]):
     def __pow__(self, __other: _E) -> _T_co: ...
 
+class _SupportsPow3NoneOnly(Protocol[_E, _T_co]):
+    def __pow__(self, __other: _E, __modulo: None = ...) -> _T_co: ...
+
 class _SupportsPow3(Protocol[_E, _M, _T_co]):
     def __pow__(self, __other: _E, __modulo: _M) -> _T_co: ...
 
+_SupportsSomeKindOfPow = Union[_SupportsPow2[Any, Any], _SupportsPow3NoneOnly[Any, Any], _SupportsPow3[Any, Any, Any]]
+
 if sys.version_info >= (3, 8):
     @overload
     def pow(base: int, exp: int, mod: Literal[0]) -> NoReturn: ...
+    @overload
+    def pow(base: int, exp: int, mod: int) -> int: ...
+    @overload
+    def pow(base: int, exp: _PositiveInteger, mod: None = ...) -> int: ...  # type: ignore[misc]
+    @overload
+    def pow(base: int, exp: _NegativeInteger, mod: None = ...) -> float: ...  # type: ignore[misc]
     # int base & positive-int exp -> int; int base & negative-int exp -> float
     # return type must be Any as `int | float` causes too many false-positive errors
     @overload
     def pow(base: int, exp: int, mod: None = ...) -> Any: ...
     @overload
-    def pow(base: int, exp: int, mod: int) -> int: ...
-    @overload
     def pow(base: float, exp: int, mod: None = ...) -> float: ...
     # float base & float exp could return float or complex
-    # return type must be Any as `float | complex` causes too many false-positive errors
+    # return type must be Any (same as complex base, complex exp),
+    # as `float | complex` causes too many false-positive errors
+    @overload
+    def pow(base: float, exp: complex | _SupportsSomeKindOfPow, mod: None = ...) -> Any: ...
     @overload
-    def pow(base: float, exp: float, mod: None = ...) -> Any: ...
+    def pow(base: complex, exp: complex | _SupportsSomeKindOfPow, mod: None = ...) -> complex: ...
     @overload
-    def pow(base: _SupportsPow2[_E, _T_co], exp: _E) -> _T_co: ...
+    def pow(base: _SupportsPow2[_E, _T_co], exp: _E, mod: None = ...) -> _T_co: ...
     @overload
-    def pow(base: _SupportsPow3[_E, _M, _T_co], exp: _E, mod: _M) -> _T_co: ...
+    def pow(base: _SupportsPow3NoneOnly[_E, _T_co], exp: _E, mod: None = ...) -> _T_co: ...
+    @overload
+    def pow(base: _SupportsPow3[_E, _M, _T_co], exp: _E, mod: _M = ...) -> _T_co: ...
+    @overload
+    def pow(base: _SupportsSomeKindOfPow, exp: float, mod: None = ...) -> Any: ...
+    @overload
+    def pow(base: _SupportsSomeKindOfPow, exp: complex, mod: None = ...) -> complex: ...
 
 else:
     @overload
     def pow(__base: int, __exp: int, __mod: Literal[0]) -> NoReturn: ...
     @overload
-    def pow(__base: int, __exp: int, __mod: None = ...) -> Any: ...
-    @overload
     def pow(__base: int, __exp: int, __mod: int) -> int: ...
     @overload
+    def pow(__base: int, __exp: _PositiveInteger, __mod: None = ...) -> int: ...  # type: ignore[misc]
+    @overload
+    def pow(__base: int, __exp: _NegativeInteger, __mod: None = ...) -> float: ...  # type: ignore[misc]
+    @overload
+    def pow(__base: int, __exp: int, __mod: None = ...) -> Any: ...
+    @overload
     def pow(__base: float, __exp: int, __mod: None = ...) -> float: ...
     @overload
-    def pow(__base: float, __exp: float, __mod: None = ...) -> Any: ...
+    def pow(__base: float, __exp: complex | _SupportsSomeKindOfPow, __mod: None = ...) -> Any: ...
+    @overload
+    def pow(__base: complex, __exp: complex | _SupportsSomeKindOfPow, __mod: None = ...) -> complex: ...
+    @overload
+    def pow(__base: _SupportsPow2[_E, _T_co], __exp: _E, __mod: None = ...) -> _T_co: ...
+    @overload
+    def pow(__base: _SupportsPow3NoneOnly[_E, _T_co], __exp: _E, __mod: None = ...) -> _T_co: ...
+    @overload
+    def pow(__base: _SupportsPow3[_E, _M, _T_co], __exp: _E, __mod: _M = ...) -> _T_co: ...
     @overload
-    def pow(__base: _SupportsPow2[_E, _T_co], __exp: _E) -> _T_co: ...
+    def pow(__base: _SupportsSomeKindOfPow, __exp: float, __mod: None = ...) -> Any: ...
     @overload
-    def pow(__base: _SupportsPow3[_E, _M, _T_co], __exp: _E, __mod: _M) -> _T_co: ...
+    def pow(__base: _SupportsSomeKindOfPow, __exp: complex, __mod: None = ...) -> complex: ...
 
 def quit(code: object = ...) -> NoReturn: ...