[mypyc] Fix super() method calls in class methods and generators

mypyc/irbuild/expression.py

@@ -402,21 +402,37 @@ def transform_super_expr(builder: IRBuilder, o: SuperExpr) -> Value:
     else:
         assert o.info is not None
         typ = builder.load_native_type_object(o.info.fullname)
-        ir = builder.mapper.type_to_ir[o.info]
-        iter_env = iter(builder.builder.args)
-        # Grab first argument
-        vself: Value = next(iter_env)
-        if builder.fn_info.is_generator:
-            # grab seventh argument (see comment in translate_super_method_call)
-            self_targ = list(builder.symtables[-1].values())[7]
-            vself = builder.read(self_targ, builder.fn_info.fitem.line)
-        elif not ir.is_ext_class:
-            vself = next(iter_env)  # second argument is self if non_extension class
+        vself = builder.read(builder.lookup(implicit_super_arg(builder)), o.line)
         args = [typ, vself]
     res = builder.py_call(sup_val, args, o.line)
     return builder.py_get_attr(res, o.name, o.line)
 
 
+def implicit_super_arg(builder: IRBuilder) -> Var:
+    """Return the variable that zero-argument super() uses as its second argument.
+
+    This is the first argument of the enclosing function. It isn't always the first
+    argument of the function being generated, since generators and nested functions
+    are compiled to methods of generated classes.
+    """
+    # A comprehension can have a scope of its own, but it has no arguments, and
+    # super() uses the function that contains the comprehension.
+    fn_info = next(info for info in reversed(builder.fn_infos) if not info.is_comprehension_scope)
+    return fn_info.fitem.arguments[0].variable
+
+
+def is_instance_method_self(builder: IRBuilder, var: Var) -> bool:
+    """Is this the self argument of an instance method?"""
+    if not var.is_self:
+        return False
+    # The cls argument of __new__ is also marked as a self argument.
+    for fn_info in builder.fn_infos:
+        fitem = fn_info.fitem
+        if fitem.name == "__new__" and fitem.arguments and fitem.arguments[0].variable is var:
+            return False
+    return True
+
+
 # Calls
 
 
@@ -601,6 +617,9 @@ def translate_super_method_call(builder: IRBuilder, expr: CallExpr, callee: Supe
             or callee.info is not typ_arg.node
         ):
             return translate_call(builder, expr, callee)
+        self_var = self_arg.node
+    else:
+        self_var = implicit_super_arg(builder)
 
     ir = builder.mapper.type_to_ir[callee.info]
     # Search for the method in the mro, skipping ourselves. We
@@ -631,26 +650,39 @@ def translate_super_method_call(builder: IRBuilder, expr: CallExpr, callee: Supe
         # super().prop(...) calls the property's value, so get it through super()
         return translate_call(builder, expr, callee)
 
-    arg_values = [builder.accept(arg) for arg in expr.args]
+    needs_self = decl.kind != FUNC_STATICMETHOD and decl.name != "__new__"
+    if needs_self and not (
+        is_instance_method_self(builder, self_var)
+        or (self_var.is_cls and decl.kind == FUNC_CLASSMETHOD)
+    ):
+        # We can only bind the method statically if super() is given the self argument of
+        # an instance method, or the cls argument of a class method when calling a class
+        # method. Otherwise it's an instance method looked up through the class, or the
+        # first argument of a static method (such as __new__) or of a nested function,
+        # which can be either an instance or a class.
+        return translate_call(builder, expr, callee)
+
+    arg_values: list[Value] = []
     arg_kinds, arg_names = expr.arg_kinds.copy(), expr.arg_names.copy()
 
-    if decl.kind != FUNC_STATICMETHOD and decl.name != "__new__":
-        # Grab first argument
-        vself: Value = builder.self()
-        if decl.kind == FUNC_CLASSMETHOD:
+    if needs_self:
+        # super() binds the object before the arguments are evaluated, and an argument
+        # can reassign the variable.
+        vself = builder.read(builder.lookup(self_var), expr.line)
+        if isinstance(vself, Register) and self_var in builder.reassigned_in_expr:
+            # The variable is reassigned via a walrus within the same expression, so
+            # copy the current value (see transform_name_expr).
+            temp = Register(vself.type)
+            builder.assign(temp, vself, expr.line)
+            vself = temp
+        if decl.kind == FUNC_CLASSMETHOD and not self_var.is_cls:
             vself = builder.primitive_op(type_op, [vself], expr.line)
-        elif builder.fn_info.is_generator:
-            # For generator classes, the self target is the 7th value
-            # in the symbol table (which is an ordered dict). This is sort
-            # of ugly, but we can't search by name since the 'self' parameter
-            # could be named anything, and it doesn't get added to the
-            # environment indexes.
-            self_targ = list(builder.symtables[-1].values())[7]
-            vself = builder.read(self_targ, builder.fn_info.fitem.line)
-        arg_values.insert(0, vself)
+        arg_values.append(vself)
         arg_kinds.insert(0, ARG_POS)
         arg_names.insert(0, None)
 
+    arg_values += [builder.accept(arg) for arg in expr.args]
+
     return builder.builder.call(decl, arg_values, arg_kinds, arg_names, expr.line)
 
 

mypyc/test-data/irbuild-classes.test

@@ -741,6 +741,119 @@ L0:
     r0 = T.foo(self)
     return 1
 
+[case testSuperClassMethod]
+class A:
+    @classmethod
+    def f(cls, x: int) -> int:
+        return x
+
+    def g(self) -> int:
+        return 1
+
+class B(A):
+    @classmethod
+    def f(cls, x: int) -> int:
+        return super().f(x)
+
+    def call_f(self) -> int:
+        return super().f(1)
+
+    @classmethod
+    def call_g(cls, b: B) -> int:
+        # An instance method looked up through the class isn't bound
+        return super().g(b)
+[out]
+def A.f(cls, x):
+    cls :: object
+    x :: int
+L0:
+    return x
+def A.g(self):
+    self :: __main__.A
+L0:
+    return 2
+def B.f(cls, x):
+    cls :: object
+    x, r0 :: int
+L0:
+    r0 = A.f(cls, x)
+    return r0
+def B.call_f(self):
+    self :: __main__.B
+    r0 :: object
+    r1 :: int
+L0:
+    r0 = CPy_TYPE(self)
+    r1 = A.f(r0, 2)
+    return r1
+def B.call_g(cls, b):
+    cls :: object
+    b :: __main__.B
+    r0 :: object
+    r1 :: str
+    r2, r3 :: object
+    r4 :: object[2]
+    r5 :: object_ptr
+    r6 :: object
+    r7 :: str
+    r8 :: object
+    r9 :: object[1]
+    r10 :: object_ptr
+    r11 :: object
+    r12 :: int
+L0:
+    r0 = builtins :: module
+    r1 = 'super'
+    r2 = CPyObject_GetAttr(r0, r1)
+    r3 = __main__.B :: type
+    r4 = [r3, cls]
+    r5 = load_address r4
+    r6 = PyObject_Vectorcall(r2, r5, 2, 0)
+    keep_alive r3, cls
+    r7 = 'g'
+    r8 = CPyObject_GetAttr(r6, r7)
+    r9 = [b]
+    r10 = load_address r9
+    r11 = PyObject_Vectorcall(r8, r10, 1, 0)
+    keep_alive b
+    r12 = unbox(int, r11)
+    return r12
+
+[case testSuperSelfReassignedInArgument]
+class A:
+    def f(self, x: int) -> int:
+        return x
+
+class B(A):
+    def f(self, x: int) -> int:
+        # super() binds self before the argument is evaluated
+        return super().f((self := B()).g())
+
+    def g(self) -> int:
+        return 1
+[out]
+def A.f(self, x):
+    self :: __main__.A
+    x :: int
+L0:
+    return x
+def B.f(self, x):
+    self :: __main__.B
+    x :: int
+    r0, r1 :: __main__.B
+    r2, r3 :: int
+L0:
+    r0 = self
+    r1 = B()
+    self = r1
+    r2 = r1.g()
+    r3 = A.f(r0, r2)
+    return r3
+def B.g(self):
+    self :: __main__.B
+L0:
+    return 2
+
 [case testSuperCallToObjectInitIsOmitted]
 class C:
     def __init__(self) -> None: