From e69fd44b21bcf7f2afd7a3e43b33f16c95ec72a4 Mon Sep 17 00:00:00 2001
From: David Euresti <david@euresti.com>
Date: Sun, 18 Feb 2018 06:47:11 -0800
Subject: [PATCH] Treat divmod like a binary operator

---
 mypy/nodes.py                         |  2 ++
 mypy/semanal.py                       |  6 ++++
 test-data/unit/check-expressions.test | 45 +++++++++++++++++++++++++++
 test-data/unit/fixtures/divmod.pyi    | 21 +++++++++++++
 4 files changed, 74 insertions(+)
 create mode 100644 test-data/unit/fixtures/divmod.pyi

diff --git a/mypy/nodes.py b/mypy/nodes.py
index 6375d500a8a3..e882031f8bd3 100644
--- a/mypy/nodes.py
+++ b/mypy/nodes.py
@@ -1321,6 +1321,7 @@ def accept(self, visitor: ExpressionVisitor[T]) -> T:
     '*': '__mul__',
     '/': '__truediv__',
     '%': '__mod__',
+    'divmod': '__divmod__',
     '//': '__floordiv__',
     '**': '__pow__',
     '@': '__matmul__',
@@ -1356,6 +1357,7 @@ def accept(self, visitor: ExpressionVisitor[T]) -> T:
     '__mul__': '__rmul__',
     '__truediv__': '__rtruediv__',
     '__mod__': '__rmod__',
+    '__divmod__': '__rdivmod__',
     '__floordiv__': '__rfloordiv__',
     '__pow__': '__rpow__',
     '__matmul__': '__rmatmul__',
diff --git a/mypy/semanal.py b/mypy/semanal.py
index 444838ede525..95b5c32e20ae 100644
--- a/mypy/semanal.py
+++ b/mypy/semanal.py
@@ -3272,6 +3272,12 @@ def visit_call_expr(self, expr: CallExpr) -> None:
             expr.analyzed.accept(self)
         elif refers_to_fullname(expr.callee, 'builtins.dict'):
             expr.analyzed = self.translate_dict_call(expr)
+        elif refers_to_fullname(expr.callee, 'builtins.divmod'):
+            if not self.check_fixed_args(expr, 2, 'divmod'):
+                return
+            expr.analyzed = OpExpr('divmod', expr.args[0], expr.args[1])
+            expr.analyzed.line = expr.line
+            expr.analyzed.accept(self)
         else:
             # Normal call expression.
             for a in expr.args:
diff --git a/test-data/unit/check-expressions.test b/test-data/unit/check-expressions.test
index 2d5d88e7b152..0b8607b06254 100644
--- a/test-data/unit/check-expressions.test
+++ b/test-data/unit/check-expressions.test
@@ -705,6 +705,51 @@ tmp/m.py:7: error: Invalid index type "int" for "A"; expected type "str"
 tmp/m.py:8: error: Invalid index type "int" for "A"; expected type "str"
 
 
+[case testDivmod]
+from typing import Tuple, Union, SupportsInt
+_Decimal = Union[Decimal, int]
+class Decimal(SupportsInt):
+    def __init__(self, int) -> None: ...
+    def __divmod__(self, other: _Decimal) -> Tuple[Decimal, Decimal]: ...
+    def __rdivmod__(self, other: _Decimal) -> Tuple[Decimal, Decimal]: ...
+
+i = 8
+f = 8.0
+d = Decimal(8)
+
+reveal_type(divmod(i, i))  # E: Revealed type is 'Tuple[builtins.int, builtins.int]'
+reveal_type(divmod(f, i))  # E: Revealed type is 'Tuple[builtins.float, builtins.float]'
+reveal_type(divmod(d, i))  # E: Revealed type is 'Tuple[__main__.Decimal, __main__.Decimal]'
+
+reveal_type(divmod(i, f))  # E: Revealed type is 'Tuple[builtins.float, builtins.float]'
+reveal_type(divmod(f, f))  # E: Revealed type is 'Tuple[builtins.float, builtins.float]'
+divmod(d, f)  # E: Unsupported operand types for divmod ("Decimal" and "float")
+
+reveal_type(divmod(i, d))  # E: Revealed type is 'Tuple[__main__.Decimal, __main__.Decimal]'
+divmod(f, d)  # E: Unsupported operand types for divmod ("float" and "Decimal")
+reveal_type(divmod(d, d))  # E: Revealed type is 'Tuple[__main__.Decimal, __main__.Decimal]'
+
+# Now some bad calls
+divmod()  # E: 'divmod' expects 2 arguments \
+          # E: Too few arguments for "divmod"
+divmod(7)  # E: 'divmod' expects 2 arguments \
+           # E: Too few arguments for "divmod"
+divmod(7, 8, 9)  # E: 'divmod' expects 2 arguments \
+                 # E: Too many arguments for "divmod"
+divmod(_x=7, _y=9)  # E: 'divmod' must be called with 2 positional arguments
+
+divmod('foo', 'foo')  # E: Unsupported left operand type for divmod ("str")
+divmod(i, 'foo')  # E: Unsupported operand types for divmod ("int" and "str")
+divmod(f, 'foo')  # E: Unsupported operand types for divmod ("float" and "str")
+divmod(d, 'foo')  # E: Unsupported operand types for divmod ("Decimal" and "str")
+
+divmod('foo', i)  # E: Unsupported operand types for divmod ("str" and "int")
+divmod('foo', f)  # E: Unsupported operand types for divmod ("str" and "float")
+divmod('foo', d)  # E: Unsupported operand types for divmod ("str" and "Decimal")
+
+[builtins fixtures/divmod.pyi]
+
+
 -- Unary operators
 -- ---------------
 
diff --git a/test-data/unit/fixtures/divmod.pyi b/test-data/unit/fixtures/divmod.pyi
new file mode 100644
index 000000000000..cf41c500f49b
--- /dev/null
+++ b/test-data/unit/fixtures/divmod.pyi
@@ -0,0 +1,21 @@
+from typing import TypeVar, Tuple, SupportsInt
+class object:
+    def __init__(self): pass
+
+class int(SupportsInt):
+    def __divmod__(self, other: int) -> Tuple[int, int]: pass
+    def __rdivmod__(self, other: int) -> Tuple[int, int]: pass
+
+class float(SupportsInt):
+    def __divmod__(self, other: float) -> Tuple[float, float]: pass
+    def __rdivmod__(self, other: float) -> Tuple[float, float]: pass
+
+
+class tuple: pass
+class function: pass
+class str: pass
+class type: pass
+class ellipsis: pass
+
+_N = TypeVar('_N', int, float)
+def divmod(_x: _N, _y: _N) -> Tuple[_N, _N]: ...