Add protocol overload to definition of builtins.divmod.

rchen152 · rchen152 · commit a39c3eb4fcc0 · 2024-02-13T10:03:47.000-08:00
PiperOrigin-RevId: 606466608
diff --git a/pytype/stubs/builtins/builtins.pytd b/pytype/stubs/builtins/builtins.pytd
@@ -65,9 +65,19 @@ def coerce(x: _T, y: _T) -> tuple[_T, _T]: ...  # E.g. coerce([], [1]) -> ([], [
 def compile(source, filename: str, mode: str, flags: int = ..., dont_inherit: int = ..., optimize: int = ...) -> code: ...
 def delattr(object, name: Union[str, bytes, bytearray]) -> None: ...
 def dir(*args, **kwargs) -> list[str]: ...
+
+class _SupportsDivMod(Protocol[_T, _T2, _T3]):
+  def __divmod__(self, other: _T) -> tuple[_T2, _T3]: ...
+
+@overload
 def divmod(x: int, y: int) -> tuple[int, int]: ...
+@overload
 def divmod(x: Union[int, float], y: Union[int, float]) -> tuple[float, float]: ...
+@overload
 def divmod(x: Union[int, float, complex], y: Union[int, float, complex]) -> tuple[complex, complex]: ...
+@overload
+def divmod(x: _SupportsDivMod[_T, _T2, _T3], y: _T) -> tuple[_T2, _T3]: ...
+
 def eval(src, *args, **kwargs) -> Any: ...  # Can't say *anything* about the result -- different from "-> object"
 def exec(src, *args, **kwargs) -> NoneType: ...
 def execfile(filename: str, *args, **kwargs) -> NoneType: ...
diff --git a/pytype/tests/test_builtins4.py b/pytype/tests/test_builtins4.py
@@ -768,6 +768,15 @@ def __index__(self) -> int:
       print(x[C()])
     """)
 
+  def test_divmod(self):
+    self.Check("""
+      import datetime
+      from typing import Tuple
+      assert_type(divmod(1, 2), Tuple[int, int])
+      assert_type(divmod(datetime.timedelta(1), datetime.timedelta(2)),
+                  Tuple[int, datetime.timedelta])
+    """)
+
 
 class SetMethodsTest(test_base.BaseTest):
   """Tests for methods of the `set` class."""
