[mypyc] Fix inheritance of async defs (#20044)

When inferring a precise generator return type for an async def (or
generator), make the generate returned by an override in a subclass
inherit from the base class generator. This means that the environment
has to be moved to a separate class in the base class generator.

Don't infer a precise generator return type when an override might have
a less precise return type, since it would break LSP.

Fixes mypyc/mypyc#1141.

Author: JukkaL

mypyc/codegen/emitclass.py

@@ -410,7 +410,7 @@ def setter_name(cl: ClassIR, attribute: str, names: NameGenerator) -> str:
 
 
 def generate_object_struct(cl: ClassIR, emitter: Emitter) -> None:
-    seen_attrs: set[tuple[str, RType]] = set()
+    seen_attrs: set[str] = set()
     lines: list[str] = []
     lines += ["typedef struct {", "PyObject_HEAD", "CPyVTableItem *vtable;"]
     if cl.has_method("__call__"):
@@ -427,9 +427,11 @@ def generate_object_struct(cl: ClassIR, emitter: Emitter) -> None:
                             lines.append(f"{BITMAP_TYPE} {attr};")
                             bitmap_attrs.append(attr)
             for attr, rtype in base.attributes.items():
-                if (attr, rtype) not in seen_attrs:
+                # Generated class may redefine certain attributes with different
+                # types in subclasses (this would be unsafe for user-defined classes).
+                if attr not in seen_attrs:
                     lines.append(f"{emitter.ctype_spaced(rtype)}{emitter.attr(attr)};")
-                    seen_attrs.add((attr, rtype))
+                    seen_attrs.add(attr)
 
                     if isinstance(rtype, RTuple):
                         emitter.declare_tuple_struct(rtype)

mypyc/codegen/emitmodule.py

@@ -1064,6 +1064,8 @@ def emit_module_exec_func(
                     "(PyObject *){t}_template, NULL, modname);".format(t=type_struct)
                 )
                 emitter.emit_lines(f"if (unlikely(!{type_struct}))", "    goto fail;")
+                name_prefix = cl.name_prefix(emitter.names)
+                emitter.emit_line(f"CPyDef_{name_prefix}_trait_vtable_setup();")
 
         emitter.emit_lines("if (CPyGlobalsInit() < 0)", "    goto fail;")
 

mypyc/ir/func_ir.py

@@ -149,8 +149,11 @@ def __init__(
         module_name: str,
         sig: FuncSignature,
         kind: int = FUNC_NORMAL,
+        *,
         is_prop_setter: bool = False,
         is_prop_getter: bool = False,
+        is_generator: bool = False,
+        is_coroutine: bool = False,
         implicit: bool = False,
         internal: bool = False,
     ) -> None:
@@ -161,6 +164,8 @@ def __init__(
         self.kind = kind
         self.is_prop_setter = is_prop_setter
         self.is_prop_getter = is_prop_getter
+        self.is_generator = is_generator
+        self.is_coroutine = is_coroutine
         if class_name is None:
             self.bound_sig: FuncSignature | None = None
         else:
@@ -219,6 +224,8 @@ def serialize(self) -> JsonDict:
             "kind": self.kind,
             "is_prop_setter": self.is_prop_setter,
             "is_prop_getter": self.is_prop_getter,
+            "is_generator": self.is_generator,
+            "is_coroutine": self.is_coroutine,
             "implicit": self.implicit,
             "internal": self.internal,
         }
@@ -240,10 +247,12 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> FuncDecl:
             data["module_name"],
             FuncSignature.deserialize(data["sig"], ctx),
             data["kind"],
-            data["is_prop_setter"],
-            data["is_prop_getter"],
-            data["implicit"],
-            data["internal"],
+            is_prop_setter=data["is_prop_setter"],
+            is_prop_getter=data["is_prop_getter"],
+            is_generator=data["is_generator"],
+            is_coroutine=data["is_coroutine"],
+            implicit=data["implicit"],
+            internal=data["internal"],
         )
 
 

mypyc/irbuild/context.py

@@ -98,7 +98,11 @@ def curr_env_reg(self) -> Value:
     def can_merge_generator_and_env_classes(self) -> bool:
         # In simple cases we can place the environment into the generator class,
         # instead of having two separate classes.
-        return self.is_generator and not self.is_nested and not self.contains_nested
+        if self._generator_class and not self._generator_class.ir.is_final_class:
+            result = False
+        else:
+            result = self.is_generator and not self.is_nested and not self.contains_nested
+        return result
 
 
 class ImplicitClass:

mypyc/irbuild/function.py

@@ -69,7 +69,7 @@
     instantiate_callable_class,
     setup_callable_class,
 )
-from mypyc.irbuild.context import FuncInfo
+from mypyc.irbuild.context import FuncInfo, GeneratorClass
 from mypyc.irbuild.env_class import (
     add_vars_to_env,
     finalize_env_class,
@@ -246,6 +246,12 @@ def c() -> None:
     is_generator = fn_info.is_generator
     builder.enter(fn_info, ret_type=sig.ret_type)
 
+    if is_generator:
+        fitem = builder.fn_info.fitem
+        assert isinstance(fitem, FuncDef), fitem
+        generator_class_ir = builder.mapper.fdef_to_generator[fitem]
+        builder.fn_info.generator_class = GeneratorClass(generator_class_ir)
+
     # Functions that contain nested functions need an environment class to store variables that
     # are free in their nested functions. Generator functions need an environment class to
     # store a variable denoting the next instruction to be executed when the __next__ function
@@ -357,8 +363,8 @@ def gen_func_ir(
                 builder.module_name,
                 sig,
                 func_decl.kind,
-                func_decl.is_prop_getter,
-                func_decl.is_prop_setter,
+                is_prop_getter=func_decl.is_prop_getter,
+                is_prop_setter=func_decl.is_prop_setter,
             )
             func_ir = FuncIR(func_decl, args, blocks, fitem.line, traceback_name=fitem.name)
         else:

mypyc/irbuild/generator.py

@@ -39,7 +39,7 @@
     object_rprimitive,
 )
 from mypyc.irbuild.builder import IRBuilder, calculate_arg_defaults, gen_arg_defaults
-from mypyc.irbuild.context import FuncInfo, GeneratorClass
+from mypyc.irbuild.context import FuncInfo
 from mypyc.irbuild.env_class import (
     add_args_to_env,
     add_vars_to_env,
@@ -166,10 +166,8 @@ def setup_generator_class(builder: IRBuilder) -> ClassIR:
         builder.fn_info.env_class = generator_class_ir
     else:
         generator_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_info.env_class)
-    generator_class_ir.mro = [generator_class_ir]
 
     builder.classes.append(generator_class_ir)
-    builder.fn_info.generator_class = GeneratorClass(generator_class_ir)
     return generator_class_ir
 
 

mypyc/irbuild/prepare.py

@@ -202,7 +202,15 @@ def prepare_func_def(
         else (FUNC_STATICMETHOD if fdef.is_static else FUNC_NORMAL)
     )
     sig = mapper.fdef_to_sig(fdef, options.strict_dunders_typing)
-    decl = FuncDecl(fdef.name, class_name, module_name, sig, kind)
+    decl = FuncDecl(
+        fdef.name,
+        class_name,
+        module_name,
+        sig,
+        kind,
+        is_generator=fdef.is_generator,
+        is_coroutine=fdef.is_coroutine,
+    )
     mapper.func_to_decl[fdef] = decl
     return decl
 
@@ -217,7 +225,7 @@ def create_generator_class_for_func(
     """
     assert fdef.is_coroutine or fdef.is_generator
     name = "_".join(x for x in [fdef.name, class_name] if x) + "_gen" + name_suffix
-    cir = ClassIR(name, module_name, is_generated=True, is_final_class=True)
+    cir = ClassIR(name, module_name, is_generated=True, is_final_class=class_name is None)
     cir.reuse_freed_instance = True
     mapper.fdef_to_generator[fdef] = cir
 
@@ -816,14 +824,70 @@ def adjust_generator_classes_of_methods(mapper: Mapper) -> None:
     This is a separate pass after type map has been built, since we need all classes
     to be processed to analyze class hierarchies.
     """
-    for fdef, ir in mapper.func_to_decl.items():
+
+    generator_methods = []
+
+    for fdef, fn_ir in mapper.func_to_decl.items():
         if isinstance(fdef, FuncDef) and (fdef.is_coroutine or fdef.is_generator):
-            gen_ir = create_generator_class_for_func(ir.module_name, ir.class_name, fdef, mapper)
+            gen_ir = create_generator_class_for_func(
+                fn_ir.module_name, fn_ir.class_name, fdef, mapper
+            )
             # TODO: We could probably support decorators sometimes (static and class method?)
             if not fdef.is_decorated:
-                # Give a more precise type for generators, so that we can optimize
-                # code that uses them. They return a generator object, which has a
-                # specific class. Without this, the type would have to be 'object'.
-                ir.sig.ret_type = RInstance(gen_ir)
-                if ir.bound_sig:
-                    ir.bound_sig.ret_type = RInstance(gen_ir)
+                name = fn_ir.name
+                precise_ret_type = True
+                if fn_ir.class_name is not None:
+                    class_ir = mapper.type_to_ir[fdef.info]
+                    subcls = class_ir.subclasses()
+                    if subcls is None:
+                        # Override could be of a different type, so we can't make assumptions.
+                        precise_ret_type = False
+                    else:
+                        for s in subcls:
+                            if name in s.method_decls:
+                                m = s.method_decls[name]
+                                if (
+                                    m.is_generator != fn_ir.is_generator
+                                    or m.is_coroutine != fn_ir.is_coroutine
+                                ):
+                                    # Override is of a different kind, and the optimization
+                                    # to use a precise generator return type doesn't work.
+                                    precise_ret_type = False
+                else:
+                    class_ir = None
+
+                if precise_ret_type:
+                    # Give a more precise type for generators, so that we can optimize
+                    # code that uses them. They return a generator object, which has a
+                    # specific class. Without this, the type would have to be 'object'.
+                    fn_ir.sig.ret_type = RInstance(gen_ir)
+                    if fn_ir.bound_sig:
+                        fn_ir.bound_sig.ret_type = RInstance(gen_ir)
+                    if class_ir is not None:
+                        if class_ir.is_method_final(name):
+                            gen_ir.is_final_class = True
+                        generator_methods.append((name, class_ir, gen_ir))
+
+    new_bases = {}
+
+    for name, class_ir, gen in generator_methods:
+        # For generator methods, we need to have subclass generator classes inherit from
+        # baseclass generator classes when there are overrides to maintain LSP.
+        base = class_ir.real_base()
+        if base is not None:
+            if base.has_method(name):
+                base_sig = base.method_sig(name)
+                if isinstance(base_sig.ret_type, RInstance):
+                    base_gen = base_sig.ret_type.class_ir
+                    new_bases[gen] = base_gen
+
+    # Add generator inheritance relationships by adjusting MROs.
+    for deriv, base in new_bases.items():
+        if base.children is not None:
+            base.children.append(deriv)
+        while True:
+            deriv.mro.append(base)
+            deriv.base_mro.append(base)
+            if base not in new_bases:
+                break
+            base = new_bases[base]

mypyc/test-data/run-async.test

@@ -1291,3 +1291,77 @@ class CancelledError(Exception): ...
 def run(x: object) -> object: ...
 def get_running_loop() -> Any: ...
 def create_task(x: object) -> Any: ...
+
+[case testAsyncInheritance1]
+from typing import final, Coroutine, Any, TypeVar
+
+import asyncio
+
+class Base1:
+    async def foo(self) -> int:
+        return 1
+
+class Derived1(Base1):
+    async def foo(self) -> int:
+        return await super().foo() + 1
+
+async def base1_foo(b: Base1) -> int:
+    return await b.foo()
+
+async def derived1_foo(b: Derived1) -> int:
+    return await b.foo()
+
+def test_async_inheritance() -> None:
+     assert asyncio.run(base1_foo(Base1())) == 1
+     assert asyncio.run(base1_foo(Derived1())) == 2
+     assert asyncio.run(derived1_foo(Derived1())) == 2
+
+@final
+class FinalClass:
+    async def foo(self) -> int:
+        return 3
+
+async def final_class_foo(b: FinalClass) -> int:
+    return await b.foo()
+
+def test_final_class() -> None:
+     assert asyncio.run(final_class_foo(FinalClass())) == 3
+
+class Base2:
+    async def foo(self) -> int:
+        return 4
+
+    async def bar(self) -> int:
+        return 5
+
+class Derived2(Base2):
+    # Does not override "foo"
+    async def bar(self) -> int:
+        return 6
+
+async def base2_foo(b: Base2) -> int:
+    return await b.foo()
+
+def test_no_override() -> None:
+     assert asyncio.run(base2_foo(Base2())) == 4
+     assert asyncio.run(base2_foo(Derived2())) == 4
+
+class Base3:
+    async def foo(self) -> int:
+        return 7
+
+class Derived3(Base3):
+    def foo(self) -> Coroutine[Any, Any, int]:
+        async def inner() -> int:
+            return 8
+        return inner()
+
+async def base3_foo(b: Base3) -> int:
+    return await b.foo()
+
+def test_override_non_async() -> None:
+     assert asyncio.run(base3_foo(Base3())) == 7
+     assert asyncio.run(base3_foo(Derived3())) == 8
+
+[file asyncio/__init__.pyi]
+def run(x: object) -> object: ...

mypyc/test-data/run-generators.test

@@ -907,3 +907,32 @@ def test_same_names() -> None:
         # matches the variable name in the input code, since internally it's generated
         # with a prefix.
         list(undefined())
+
+[case testGeneratorInheritance]
+from typing import Iterator
+
+class Base1:
+    def foo(self) -> Iterator[int]:
+        yield 1
+
+class Derived1(Base1):
+    def foo(self) -> Iterator[int]:
+        yield 2
+        yield 3
+
+def base1_foo(b: Base1) -> list[int]:
+    a = []
+    for x in b.foo():
+        a.append(x)
+    return a
+
+def derived1_foo(b: Derived1) -> list[int]:
+    a = []
+    for x in b.foo():
+        a.append(x)
+    return a
+
+def test_generator_override() -> None:
+     assert base1_foo(Base1()) == [1]
+     assert base1_foo(Derived1()) == [2, 3]
+     assert derived1_foo(Derived1()) == [2, 3]