[mypyc] Support __(r)divmod__ dunders (#14613)

Pretty simple. Towards mypyc/mypyc#553.

Author: ichard26

mypyc/codegen/emitclass.py

@@ -82,6 +82,8 @@ def wrapper_slot(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
     "__rtruediv__": ("nb_true_divide", generate_bin_op_wrapper),
     "__floordiv__": ("nb_floor_divide", generate_bin_op_wrapper),
     "__rfloordiv__": ("nb_floor_divide", generate_bin_op_wrapper),
+    "__divmod__": ("nb_divmod", generate_bin_op_wrapper),
+    "__rdivmod__": ("nb_divmod", generate_bin_op_wrapper),
     "__lshift__": ("nb_lshift", generate_bin_op_wrapper),
     "__rlshift__": ("nb_lshift", generate_bin_op_wrapper),
     "__rshift__": ("nb_rshift", generate_bin_op_wrapper),

mypyc/primitives/generic_ops.py

@@ -75,6 +75,17 @@
         priority=0,
     )
 
+
+function_op(
+    name="builtins.divmod",
+    arg_types=[object_rprimitive, object_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name="PyNumber_Divmod",
+    error_kind=ERR_MAGIC,
+    priority=0,
+)
+
+
 for op, funcname in [
     ("+=", "PyNumber_InPlaceAdd"),
     ("-=", "PyNumber_InPlaceSubtract"),

mypyc/test-data/fixtures/ir.py

@@ -8,13 +8,19 @@
 
 T = TypeVar('T')
 T_co = TypeVar('T_co', covariant=True)
+T_contra = TypeVar('T_contra', contravariant=True)
 S = TypeVar('S')
 K = TypeVar('K') # for keys in mapping
 V = TypeVar('V') # for values in mapping
 
 class __SupportsAbs(Protocol[T_co]):
     def __abs__(self) -> T_co: pass
 
+class __SupportsDivMod(Protocol[T_contra, T_co]):
+    def __divmod__(self, other: T_contra) -> T_co: ...
+
+class __SupportsRDivMod(Protocol[T_contra, T_co]):
+    def __rdivmod__(self, other: T_contra) -> T_co: ...
 
 class object:
     def __init__(self) -> None: pass
@@ -42,6 +48,7 @@ def __pow__(self, n: int, modulo: Optional[int] = None) -> int: pass
     def __floordiv__(self, x: int) -> int: pass
     def __truediv__(self, x: float) -> float: pass
     def __mod__(self, x: int) -> int: pass
+    def __divmod__(self, x: float) -> Tuple[float, float]: pass
     def __neg__(self) -> int: pass
     def __pos__(self) -> int: pass
     def __abs__(self) -> int: pass
@@ -307,6 +314,10 @@ def zip(x: Iterable[T], y: Iterable[S]) -> Iterator[Tuple[T, S]]: ...
 def zip(x: Iterable[T], y: Iterable[S], z: Iterable[V]) -> Iterator[Tuple[T, S, V]]: ...
 def eval(e: str) -> Any: ...
 def abs(x: __SupportsAbs[T]) -> T: ...
+@overload
+def divmod(x: __SupportsDivMod[T_contra, T_co], y: T_contra) -> T_co: ...
+@overload
+def divmod(x: T_contra, y: __SupportsRDivMod[T_contra, T_co]) -> T_co: ...
 def exit() -> None: ...
 def min(x: T, y: T) -> T: ...
 def max(x: T, y: T) -> T: ...

mypyc/test-data/irbuild-any.test

@@ -198,3 +198,18 @@ L0:
     b = r4
     return 1
 
+[case testFunctionBasedOps]
+def f() -> None:
+    a = divmod(5, 2)
+[out]
+def f():
+    r0, r1, r2 :: object
+    r3, a :: tuple[float, float]
+L0:
+    r0 = object 5
+    r1 = object 2
+    r2 = PyNumber_Divmod(r0, r1)
+    r3 = unbox(tuple[float, float], r2)
+    a = r3
+    return 1
+

mypyc/test-data/run-dunders.test

@@ -402,6 +402,9 @@ class C:
     def __floordiv__(self, y: int) -> int:
         return self.x + y + 30
 
+    def __divmod__(self, y: int) -> int:
+        return self.x + y + 40
+
 def test_generic() -> None:
     a: Any = C()
     assert a + 3 == 8
@@ -417,11 +420,13 @@ def test_generic() -> None:
     assert a @ 3 == 18
     assert a / 2 == 27
     assert a // 2 == 37
+    assert divmod(a, 2) == 47
 
 def test_native() -> None:
     c = C()
     assert c + 3 == 8
     assert c - 3 == 2
+    assert divmod(c, 3) == 48
 
 def test_error() -> None:
     a: Any = C()