diff --git a/mypyc/codegen/emitclass.py b/mypyc/codegen/emitclass.py index 122f62a0d582..d64940084f12 100644 --- a/mypyc/codegen/emitclass.py +++ b/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) diff --git a/mypyc/codegen/emitmodule.py b/mypyc/codegen/emitmodule.py index 3602b3c26e03..7ec315a6bd34 100644 --- a/mypyc/codegen/emitmodule.py +++ b/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;") diff --git a/mypyc/ir/func_ir.py b/mypyc/ir/func_ir.py index 881ac5939c27..d11fef42feb5 100644 --- a/mypyc/ir/func_ir.py +++ b/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"], ) diff --git a/mypyc/irbuild/context.py b/mypyc/irbuild/context.py index 8d2e55ed96fb..d5a48bf838c8 100644 --- a/mypyc/irbuild/context.py +++ b/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: diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py index c9f999597d30..738d19ea6748 100644 --- a/mypyc/irbuild/function.py +++ b/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: diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index b3a417ed6a3e..4dcd748f6eff 100644 --- a/mypyc/irbuild/generator.py +++ b/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 diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index 2d0a1a8f03bf..0f7cc7e3b3c5 100644 --- a/mypyc/irbuild/prepare.py +++ b/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] diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test index 94a1cd2e97c5..cf063310fd89 100644 --- a/mypyc/test-data/run-async.test +++ b/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: ... diff --git a/mypyc/test-data/run-generators.test b/mypyc/test-data/run-generators.test index bfbd5b83696b..c8e83173474d 100644 --- a/mypyc/test-data/run-generators.test +++ b/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]