From d92c4376497fe1ce81621ee3a403e65b8236f60b Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Mon, 9 Jun 2025 16:26:34 +0100 Subject: [PATCH 01/22] WIP populate generator classes before main irbuild phase Also use generator class as the return type of generators. --- mypyc/irbuild/generator.py | 6 ++++-- mypyc/irbuild/mapper.py | 6 +++++- mypyc/irbuild/prepare.py | 5 +++++ 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 0e4b0e3e184a8..896d015409078 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -155,8 +155,10 @@ def instantiate_generator_class(builder: IRBuilder) -> Value: def setup_generator_class(builder: IRBuilder) -> ClassIR: name = f"{builder.fn_info.namespaced_name()}_gen" - generator_class_ir = ClassIR(name, builder.module_name, is_generated=True, is_final_class=True) - generator_class_ir.reuse_freed_instance = True + m = builder.mapper + generator_class_ir = m.fdef_to_generator[builder.fn_info.fitem] + generator_class_ir.name = name + generator_class_ir.ctor.name = name if builder.fn_info.can_merge_generator_and_env_classes(): builder.fn_info.env_class = generator_class_ir else: diff --git a/mypyc/irbuild/mapper.py b/mypyc/irbuild/mapper.py index 7c6e03d0037cd..4ab15d99d7c7d 100644 --- a/mypyc/irbuild/mapper.py +++ b/mypyc/irbuild/mapper.py @@ -64,6 +64,7 @@ def __init__(self, group_map: dict[str, str | None]) -> None: self.type_to_ir: dict[TypeInfo, ClassIR] = {} self.func_to_decl: dict[SymbolNode, FuncDecl] = {} self.symbol_fullnames: set[str] = set() + self.fdef_to_generator: dict[FuncDef, ClassIR] = {} def type_to_rtype(self, typ: Type | None) -> RType: if typ is None: @@ -171,7 +172,10 @@ def fdef_to_sig(self, fdef: FuncDef, strict_dunders_typing: bool) -> FuncSignatu for typ, kind in zip(fdef.type.arg_types, fdef.type.arg_kinds) ] arg_pos_onlys = [name is None for name in fdef.type.arg_names] - ret = self.type_to_rtype(fdef.type.ret_type) + if fdef.is_coroutine or fdef.is_generator: + ret = RInstance(self.fdef_to_generator[fdef]) + else: + ret = self.type_to_rtype(fdef.type.ret_type) else: # Handle unannotated functions arg_types = [object_rprimitive for _ in fdef.arguments] diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index 65951999dcf99..e3f546e8d9e11 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -179,6 +179,11 @@ def prepare_func_def( mapper: Mapper, options: CompilerOptions, ) -> FuncDecl: + if fdef.is_coroutine or fdef.is_generator: + cir = ClassIR(fdef.name, module_name, is_generated=True) + cir.reuse_freed_instance = True + mapper.fdef_to_generator[fdef] = cir + kind = ( FUNC_STATICMETHOD if fdef.is_static From 4c44b4d2b79189577d8afc330de53437ba11a09b Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Mon, 9 Jun 2025 16:31:17 +0100 Subject: [PATCH 02/22] Omit CPy_GetCoro call when calling native coro --- mypyc/irbuild/statement.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 5c32d8f1a50c4..322b0ac31220f 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -53,6 +53,7 @@ Assign, BasicBlock, Branch, + Call, InitStatic, Integer, LoadAddress, @@ -930,11 +931,15 @@ def emit_yield_from_or_await( to_yield_reg = Register(object_rprimitive) received_reg = Register(object_rprimitive) - get_op = coro_op if is_await else iter_op - if isinstance(get_op, PrimitiveDescription): - iter_val = builder.primitive_op(get_op, [val], line) + if isinstance(val, Call) and isinstance(val.type, RInstance) and val.type.class_ir.is_generated: + # XXX is_generator not is_generated !!!!!!! FIXME FIXME + iter_val = val else: - iter_val = builder.call_c(get_op, [val], line) + get_op = coro_op if is_await else iter_op + if isinstance(get_op, PrimitiveDescription): + iter_val = builder.primitive_op(get_op, [val], line) + else: + iter_val = builder.call_c(get_op, [val], line) iter_reg = builder.maybe_spill_assignable(iter_val) From ab448523bdc15928329b92d2c65ead88eb6e4a0d Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Mon, 9 Jun 2025 16:37:13 +0100 Subject: [PATCH 03/22] WIP better --- mypyc/irbuild/statement.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 322b0ac31220f..51f3aac7a3b1d 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -933,7 +933,9 @@ def emit_yield_from_or_await( if isinstance(val, Call) and isinstance(val.type, RInstance) and val.type.class_ir.is_generated: # XXX is_generator not is_generated !!!!!!! FIXME FIXME - iter_val = val + r = Register(val.type) + builder.assign(r, val, line) + iter_val = r else: get_op = coro_op if is_await else iter_op if isinstance(get_op, PrimitiveDescription): From 74db611e8b715cd21ecb9b9a3a1a6112defa6d99 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Mon, 9 Jun 2025 17:08:26 +0100 Subject: [PATCH 04/22] WIP attempt to fix naming --- mypyc/irbuild/generator.py | 6 +++--- mypyc/irbuild/prepare.py | 3 ++- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 896d015409078..73c2975ea079d 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -153,12 +153,12 @@ def instantiate_generator_class(builder: IRBuilder) -> Value: def setup_generator_class(builder: IRBuilder) -> ClassIR: - name = f"{builder.fn_info.namespaced_name()}_gen" + #name = f"{builder.fn_info.namespaced_name()}_gen" m = builder.mapper generator_class_ir = m.fdef_to_generator[builder.fn_info.fitem] - generator_class_ir.name = name - generator_class_ir.ctor.name = name + #generator_class_ir.name = name + #generator_class_ir.ctor.name = name if builder.fn_info.can_merge_generator_and_env_classes(): builder.fn_info.env_class = generator_class_ir else: diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index e3f546e8d9e11..bf743240ffcb8 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -180,7 +180,8 @@ def prepare_func_def( options: CompilerOptions, ) -> FuncDecl: if fdef.is_coroutine or fdef.is_generator: - cir = ClassIR(fdef.name, module_name, is_generated=True) + name = "_".join(x for x in [fdef.name, class_name] if x) + "_gen" + cir = ClassIR(name, module_name, is_generated=True) cir.reuse_freed_instance = True mapper.fdef_to_generator[fdef] = cir From 9959e24b16693c854160e4429ed30aa61f6ec630 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Tue, 10 Jun 2025 11:16:07 +0100 Subject: [PATCH 05/22] Fix segfault --- mypyc/irbuild/statement.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 51f3aac7a3b1d..04e1e2427dd8c 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -955,7 +955,7 @@ def emit_yield_from_or_await( builder.assign(result, builder.call_c(check_stop_op, [], line), line) # Clear the spilled iterator/coroutine so that it will be freed. # Otherwise, the freeing of the spilled register would likely be delayed. - err = builder.add(LoadErrorValue(object_rprimitive)) + err = builder.add(LoadErrorValue(iter_reg.type)) builder.assign(iter_reg, err, line) builder.goto(done_block) From baca816f8ae6d25bc53ad88af1e0229a66f20411 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Tue, 10 Jun 2025 11:20:25 +0100 Subject: [PATCH 06/22] Simplify --- mypyc/irbuild/statement.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 04e1e2427dd8c..e1df429b21d40 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -933,9 +933,7 @@ def emit_yield_from_or_await( if isinstance(val, Call) and isinstance(val.type, RInstance) and val.type.class_ir.is_generated: # XXX is_generator not is_generated !!!!!!! FIXME FIXME - r = Register(val.type) - builder.assign(r, val, line) - iter_val = r + iter_val = val else: get_op = coro_op if is_await else iter_op if isinstance(get_op, PrimitiveDescription): From e1282a75ff655fc6c96e3b6ce1c4b20fe9da2466 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Tue, 10 Jun 2025 11:41:52 +0100 Subject: [PATCH 07/22] Call __next__ directly if possible in await --- mypyc/irbuild/statement.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index e1df429b21d40..b18f6ed7674ad 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -48,6 +48,7 @@ ) from mypyc.common import TEMP_ATTR_NAME from mypyc.ir.ops import ( + ERR_NEVER, NAMESPACE_MODULE, NO_TRACEBACK_LINE_NO, Assign, @@ -944,7 +945,14 @@ def emit_yield_from_or_await( iter_reg = builder.maybe_spill_assignable(iter_val) stop_block, main_block, done_block = BasicBlock(), BasicBlock(), BasicBlock() - _y_init = builder.call_c(next_raw_op, [builder.read(iter_reg)], line) + + if isinstance(iter_reg.type, RInstance) and iter_reg.type.class_ir.has_method("__next__"): + m = MethodCall(builder.read(iter_reg), "__next__", [], line) + m.error_kind = ERR_NEVER + _y_init = builder.add(m) + else: + _y_init = builder.call_c(next_raw_op, [builder.read(iter_reg)], line) + builder.add(Branch(_y_init, stop_block, main_block, Branch.IS_ERROR)) # Try extracting a return value from a StopIteration and return it. From 5e6b575d52f29903fa8d4fef2b668059795dbfd8 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Tue, 10 Jun 2025 13:22:05 +0100 Subject: [PATCH 08/22] Call helper method directly --- mypyc/irbuild/generator.py | 18 +----------------- mypyc/irbuild/prepare.py | 23 ++++++++++++++++++++++- mypyc/irbuild/statement.py | 7 +++++-- 3 files changed, 28 insertions(+), 20 deletions(-) diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 73c2975ea079d..1a60e65db2dc1 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -241,23 +241,7 @@ def add_helper_to_generator_class( fn_info: FuncInfo, ) -> FuncDecl: """Generates a helper method for a generator class, called by '__next__' and 'throw'.""" - sig = FuncSignature( - ( - RuntimeArg(SELF_NAME, object_rprimitive), - RuntimeArg("type", object_rprimitive), - RuntimeArg("value", object_rprimitive), - RuntimeArg("traceback", object_rprimitive), - RuntimeArg("arg", object_rprimitive), - ), - sig.ret_type, - ) - helper_fn_decl = FuncDecl( - "__mypyc_generator_helper__", - fn_info.generator_class.ir.name, - builder.module_name, - sig, - internal=True, - ) + helper_fn_decl = fn_info.generator_class.ir.method_decls["__mypyc_generator_helper__"] helper_fn_ir = FuncIR( helper_fn_decl, arg_regs, blocks, fn_info.fitem.line, traceback_name=fn_info.fitem.name ) diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index bf743240ffcb8..c6caa75172caf 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -38,7 +38,7 @@ from mypy.semanal import refers_to_fullname from mypy.traverser import TraverserVisitor from mypy.types import Instance, Type, get_proper_type -from mypyc.common import PROPSET_PREFIX, get_id_from_name +from mypyc.common import PROPSET_PREFIX, get_id_from_name, SELF_NAME from mypyc.crash import catch_errors from mypyc.errors import Errors from mypyc.ir.class_ir import ClassIR @@ -62,6 +62,7 @@ ) from mypyc.options import CompilerOptions from mypyc.sametype import is_same_type +from mypyc.ir.rtypes import object_rprimitive def build_type_map( @@ -185,6 +186,26 @@ def prepare_func_def( cir.reuse_freed_instance = True mapper.fdef_to_generator[fdef] = cir + helper_sig = FuncSignature( + ( + RuntimeArg(SELF_NAME, object_rprimitive), + RuntimeArg("type", object_rprimitive), + RuntimeArg("value", object_rprimitive), + RuntimeArg("traceback", object_rprimitive), + RuntimeArg("arg", object_rprimitive), + ), + RInstance(cir), + ) + + helper_fn_decl = FuncDecl( + "__mypyc_generator_helper__", + name, + module_name, + helper_sig, + internal=True, + ) + cir.method_decls[helper_fn_decl.name] = helper_fn_decl + kind = ( FUNC_STATICMETHOD if fdef.is_static diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index b18f6ed7674ad..1d5be5385dd4c 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -946,8 +946,11 @@ def emit_yield_from_or_await( stop_block, main_block, done_block = BasicBlock(), BasicBlock(), BasicBlock() - if isinstance(iter_reg.type, RInstance) and iter_reg.type.class_ir.has_method("__next__"): - m = MethodCall(builder.read(iter_reg), "__next__", [], line) + helper_method = "__mypyc_generator_helper__" + if isinstance(iter_reg.type, RInstance) and iter_reg.type.class_ir.has_method(helper_method): + obj = builder.read(iter_reg) + nn = builder.none_object() + m = MethodCall(obj, helper_method, [nn, nn, nn, nn], line) m.error_kind = ERR_NEVER _y_init = builder.add(m) else: From e3d3578976b9291b2d94d0fa4747a29a1b1835c5 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Thu, 19 Jun 2025 12:52:45 +0100 Subject: [PATCH 09/22] Generator class is final (fix after rebase) --- mypyc/irbuild/prepare.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index c6caa75172caf..93587e64a1ac1 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -182,7 +182,7 @@ def prepare_func_def( ) -> FuncDecl: if fdef.is_coroutine or fdef.is_generator: name = "_".join(x for x in [fdef.name, class_name] if x) + "_gen" - cir = ClassIR(name, module_name, is_generated=True) + cir = ClassIR(name, module_name, is_generated=True, is_final_class=True) cir.reuse_freed_instance = True mapper.fdef_to_generator[fdef] = cir From 7bdc1f54dcf9eef5b91954535ec1fa0f602b9122 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 20 Jun 2025 15:55:30 +0100 Subject: [PATCH 10/22] Fix helper method return type -- it should be object --- mypyc/irbuild/function.py | 2 +- mypyc/irbuild/generator.py | 30 ++++++++++++------------------ mypyc/irbuild/prepare.py | 2 +- 3 files changed, 14 insertions(+), 20 deletions(-) diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py index dbebc350bb6c1..dcc5a306bcde5 100644 --- a/mypyc/irbuild/function.py +++ b/mypyc/irbuild/function.py @@ -261,7 +261,7 @@ def c() -> None: ) # Re-enter the FuncItem and visit the body of the function this time. - gen_generator_func_body(builder, fn_info, sig, func_reg) + gen_generator_func_body(builder, fn_info, func_reg) else: func_ir, func_reg = gen_func_body(builder, sig, cdef, is_singledispatch) diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 1a60e65db2dc1..188c3982fdf25 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -79,7 +79,7 @@ def gen_generator_func( def gen_generator_func_body( - builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature, func_reg: Value | None + builder: IRBuilder, fn_info: FuncInfo, func_reg: Value | None ) -> None: """Generate IR based on the body of a generator function. @@ -88,7 +88,7 @@ class that implements the function (each function gets a separate class). Return the symbol table for the body. """ - builder.enter(fn_info, ret_type=sig.ret_type) + builder.enter(fn_info, ret_type=object_rprimitive) setup_env_for_generator_class(builder) load_outer_envs(builder, builder.fn_info.generator_class) @@ -117,7 +117,7 @@ class that implements the function (each function gets a separate class). args, _, blocks, ret_type, fn_info = builder.leave() - add_methods_to_generator_class(builder, fn_info, sig, args, blocks, fitem.is_coroutine) + add_methods_to_generator_class(builder, fn_info, args, blocks, fitem.is_coroutine) # Evaluate argument defaults in the surrounding scope, since we # calculate them *once* when the function definition is evaluated. @@ -153,12 +153,8 @@ def instantiate_generator_class(builder: IRBuilder) -> Value: def setup_generator_class(builder: IRBuilder) -> ClassIR: - #name = f"{builder.fn_info.namespaced_name()}_gen" - - m = builder.mapper - generator_class_ir = m.fdef_to_generator[builder.fn_info.fitem] - #generator_class_ir.name = name - #generator_class_ir.ctor.name = name + mapper = builder.mapper + generator_class_ir = mapper.fdef_to_generator[builder.fn_info.fitem] if builder.fn_info.can_merge_generator_and_env_classes(): builder.fn_info.env_class = generator_class_ir else: @@ -218,16 +214,15 @@ def add_raise_exception_blocks_to_generator_class(builder: IRBuilder, line: int) def add_methods_to_generator_class( builder: IRBuilder, fn_info: FuncInfo, - sig: FuncSignature, arg_regs: list[Register], blocks: list[BasicBlock], is_coroutine: bool, ) -> None: - helper_fn_decl = add_helper_to_generator_class(builder, arg_regs, blocks, sig, fn_info) - add_next_to_generator_class(builder, fn_info, helper_fn_decl, sig) - add_send_to_generator_class(builder, fn_info, helper_fn_decl, sig) + helper_fn_decl = add_helper_to_generator_class(builder, arg_regs, blocks, fn_info) + add_next_to_generator_class(builder, fn_info, helper_fn_decl) + add_send_to_generator_class(builder, fn_info, helper_fn_decl) add_iter_to_generator_class(builder, fn_info) - add_throw_to_generator_class(builder, fn_info, helper_fn_decl, sig) + add_throw_to_generator_class(builder, fn_info, helper_fn_decl) add_close_to_generator_class(builder, fn_info) if is_coroutine: add_await_to_generator_class(builder, fn_info) @@ -237,7 +232,6 @@ def add_helper_to_generator_class( builder: IRBuilder, arg_regs: list[Register], blocks: list[BasicBlock], - sig: FuncSignature, fn_info: FuncInfo, ) -> FuncDecl: """Generates a helper method for a generator class, called by '__next__' and 'throw'.""" @@ -259,7 +253,7 @@ def add_iter_to_generator_class(builder: IRBuilder, fn_info: FuncInfo) -> None: def add_next_to_generator_class( - builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl, sig: FuncSignature + builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl ) -> None: """Generates the '__next__' method for a generator class.""" with builder.enter_method(fn_info.generator_class.ir, "__next__", object_rprimitive, fn_info): @@ -276,7 +270,7 @@ def add_next_to_generator_class( def add_send_to_generator_class( - builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl, sig: FuncSignature + builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl ) -> None: """Generates the 'send' method for a generator class.""" with builder.enter_method(fn_info.generator_class.ir, "send", object_rprimitive, fn_info): @@ -294,7 +288,7 @@ def add_send_to_generator_class( def add_throw_to_generator_class( - builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl, sig: FuncSignature + builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl ) -> None: """Generates the 'throw' method for a generator class.""" with builder.enter_method(fn_info.generator_class.ir, "throw", object_rprimitive, fn_info): diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index 93587e64a1ac1..e0ec2c9760a6c 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -194,7 +194,7 @@ def prepare_func_def( RuntimeArg("traceback", object_rprimitive), RuntimeArg("arg", object_rprimitive), ), - RInstance(cir), + object_rprimitive, ) helper_fn_decl = FuncDecl( From a925378bb0f0f257df1b64afe187de6c0f373472 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 20 Jun 2025 16:36:09 +0100 Subject: [PATCH 11/22] Don't apply to decorated functions, since we don't know concrete return type --- mypyc/irbuild/mapper.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/mapper.py b/mypyc/irbuild/mapper.py index 4ab15d99d7c7d..0fd59c6c7506c 100644 --- a/mypyc/irbuild/mapper.py +++ b/mypyc/irbuild/mapper.py @@ -172,7 +172,8 @@ def fdef_to_sig(self, fdef: FuncDef, strict_dunders_typing: bool) -> FuncSignatu for typ, kind in zip(fdef.type.arg_types, fdef.type.arg_kinds) ] arg_pos_onlys = [name is None for name in fdef.type.arg_names] - if fdef.is_coroutine or fdef.is_generator: + # TODO: We could probably support decorators sometimes (static and class method?) + if (fdef.is_coroutine or fdef.is_generator) and not fdef.is_decorated: ret = RInstance(self.fdef_to_generator[fdef]) else: ret = self.type_to_rtype(fdef.type.ret_type) From 3328d18cf956ea3ed266cbdcaf92754eff42b0eb Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Wed, 2 Jul 2025 17:03:22 +0100 Subject: [PATCH 12/22] Fix nested generators --- mypyc/irbuild/main.py | 11 +++++-- mypyc/irbuild/prepare.py | 30 ++++++++++-------- mypyc/test-data/run-async.test | 35 ++++++++++++++++++++- mypyc/test-data/run-generators.test | 48 +++++++++++++++++++++++------ 4 files changed, 100 insertions(+), 24 deletions(-) diff --git a/mypyc/irbuild/main.py b/mypyc/irbuild/main.py index 7cdc6b6867785..eda5c0b0861b9 100644 --- a/mypyc/irbuild/main.py +++ b/mypyc/irbuild/main.py @@ -25,7 +25,7 @@ def f(x: int) -> int: from typing import Any, Callable, TypeVar, cast from mypy.build import Graph -from mypy.nodes import ClassDef, Expression, MypyFile +from mypy.nodes import ClassDef, Expression, MypyFile, FuncDef from mypy.state import state from mypy.types import Type from mypyc.analysis.attrdefined import analyze_always_defined_attrs @@ -37,7 +37,7 @@ def f(x: int) -> int: from mypyc.irbuild.builder import IRBuilder from mypyc.irbuild.mapper import Mapper from mypyc.irbuild.prebuildvisitor import PreBuildVisitor -from mypyc.irbuild.prepare import build_type_map, find_singledispatch_register_impls +from mypyc.irbuild.prepare import build_type_map, find_singledispatch_register_impls, create_generator_class_if_needed from mypyc.irbuild.visitor import IRBuilderVisitor from mypyc.irbuild.vtable import compute_vtable from mypyc.options import CompilerOptions @@ -76,6 +76,13 @@ def build_ir( pbv = PreBuildVisitor(errors, module, singledispatch_info.decorators_to_remove, types) module.accept(pbv) + # Declare generator classes for nested async functions and generators. + for fdef in pbv.nested_funcs: + if isinstance(fdef, FuncDef): + # Make generator class name sufficiently unique. + suffix = f"___{fdef.line}" + create_generator_class_if_needed(module.fullname, None, fdef, mapper, name_suffix=suffix) + # Construct and configure builder objects (cyclic runtime dependency). visitor = IRBuilderVisitor() builder = IRBuilder( diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index e0ec2c9760a6c..82be3dc91bdb6 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -116,7 +116,7 @@ def build_type_map( # Collect all the functions also. We collect from the symbol table # so that we can easily pick out the right copy of a function that - # is conditionally defined. + # is conditionally defined. This doesn't include nested functions! for module in modules: for func in get_module_func_defs(module): prepare_func_def(module.fullname, None, func, mapper, options) @@ -180,8 +180,24 @@ def prepare_func_def( mapper: Mapper, options: CompilerOptions, ) -> FuncDecl: + create_generator_class_if_needed(module_name, class_name, fdef, mapper) + + kind = ( + FUNC_STATICMETHOD + if fdef.is_static + else (FUNC_CLASSMETHOD if fdef.is_class else FUNC_NORMAL) + ) + sig = mapper.fdef_to_sig(fdef, options.strict_dunders_typing) + decl = FuncDecl(fdef.name, class_name, module_name, sig, kind) + mapper.func_to_decl[fdef] = decl + return decl + + +def create_generator_class_if_needed( + module_name: str, class_name: str | None, fdef: FuncDef, mapper: Mapper, name_suffix: str = "" +) -> None: if fdef.is_coroutine or fdef.is_generator: - name = "_".join(x for x in [fdef.name, class_name] if x) + "_gen" + 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.reuse_freed_instance = True mapper.fdef_to_generator[fdef] = cir @@ -206,16 +222,6 @@ def prepare_func_def( ) cir.method_decls[helper_fn_decl.name] = helper_fn_decl - kind = ( - FUNC_STATICMETHOD - if fdef.is_static - else (FUNC_CLASSMETHOD if fdef.is_class else FUNC_NORMAL) - ) - sig = mapper.fdef_to_sig(fdef, options.strict_dunders_typing) - decl = FuncDecl(fdef.name, class_name, module_name, sig, kind) - mapper.func_to_decl[fdef] = decl - return decl - def prepare_method_def( ir: ClassIR, diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test index 2dad720f99cd9..12c11b2fe480f 100644 --- a/mypyc/test-data/run-async.test +++ b/mypyc/test-data/run-async.test @@ -563,8 +563,10 @@ def test_bool() -> None: def run(x: object) -> object: ... [case testRunAsyncNestedFunctions] +from __future__ import annotations + import asyncio -from typing import cast, Iterator +from typing import cast, Iterator, overload, Awaitable, Any, TypeVar from testutil import assertRaises @@ -641,6 +643,37 @@ def test_async_def_contains_two_nested_functions() -> None: (5 + 3 + 1), (7 + 4 + 10 + 1) ) +async def async_def_contains_overloaded_async_def(n: int) -> int: + @overload + async def f(x: int) -> int: ... + + @overload + async def f(x: str) -> str: ... + + async def f(x: int | str) -> Any: + return x + + return (await f(n)) + 1 + + +def test_async_def_contains_overloaded_async_def() -> None: + assert asyncio.run(async_def_contains_overloaded_async_def(5)) == 6 + +T = TypeVar("T") + +def deco(f: T) -> T: + return f + +async def async_def_contains_decorated_async_def(n: int) -> int: + @deco + async def f(x: int) -> int: + return x + 2 + + return (await f(n)) + 1 + + +def test_async_def_contains_decorated_async_def() -> None: + assert asyncio.run(async_def_contains_decorated_async_def(7)) == 10 [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 9c0b51d58e79f..a43aff27dd45e 100644 --- a/mypyc/test-data/run-generators.test +++ b/mypyc/test-data/run-generators.test @@ -190,7 +190,9 @@ exit! a exception! ((1,), 'exception!') [case testYieldNested] -from typing import Callable, Generator +from typing import Callable, Generator, Iterator, TypeVar, overload + +from testutil import run_generator def normal(a: int, b: float) -> Callable: def generator(x: int, y: str) -> Generator: @@ -235,15 +237,43 @@ def outer() -> Generator: yield i return recursive(10) -[file driver.py] -from native import normal, generator, triple, another_triple, outer -from testutil import run_generator +def test_return_nested_generator() -> None: + assert run_generator(normal(1, 2.0)(3, '4.00')) == ((1, 2.0, 3, '4.00'), None) + assert run_generator(generator(1)) == ((1, 2, 3), None) + assert run_generator(triple()()) == ((1, 2, 3), None) + assert run_generator(another_triple()()) == ((1,), None) + assert run_generator(outer()) == ((0, 1, 2, 3, 4), None) + +def call_nested(x: int) -> list[int]: + def generator() -> Iterator[int]: + n = int() + 2 + yield x + yield n * x + + a = [] + for x in generator(): + a.append(x) + return a + +T = TypeVar("T") + +def deco(f: T) -> T: + return f + +def call_nested_decorated(x: int) -> list[int]: + @deco + def generator() -> Iterator[int]: + n = int() + 3 + yield x + yield n * x + + a = [] + for x in generator(): + a.append(x) + return a -assert run_generator(normal(1, 2.0)(3, '4.00')) == ((1, 2.0, 3, '4.00'), None) -assert run_generator(generator(1)) == ((1, 2, 3), None) -assert run_generator(triple()()) == ((1, 2, 3), None) -assert run_generator(another_triple()()) == ((1,), None) -assert run_generator(outer()) == ((0, 1, 2, 3, 4), None) +def test_call_nested_generator_in_function() -> None: + assert call_nested_decorated(5) == [5, 15] [case testYieldThrow] from typing import Generator, Iterable, Any, Union From f29315045930df23e837fd45c96920dd26680004 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Thu, 3 Jul 2025 15:30:43 +0100 Subject: [PATCH 13/22] Add test case --- mypyc/test-data/run-async.test | 35 ++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test index 12c11b2fe480f..99276688465f3 100644 --- a/mypyc/test-data/run-async.test +++ b/mypyc/test-data/run-async.test @@ -2,6 +2,7 @@ [case testRunAsyncBasics] import asyncio +from typing import Callable, Awaitable from testutil import assertRaises @@ -72,6 +73,40 @@ def test_exception() -> None: assert asyncio.run(exc5()) == 3 assert asyncio.run(exc6()) == 3 +async def indirect_call(x: int, c: Callable[[int], Awaitable[int]]) -> int: + return await c(x) + +async def indirect_call_2(a: Awaitable[None]) -> None: + await a + +async def indirect_call_3(a: Awaitable[float]) -> float: + return (await a) + 1.0 + +async def inc(x: int) -> int: + await asyncio.sleep(0) + return x + 1 + +async def ident(x: float, err: bool = False) -> float: + await asyncio.sleep(0.0) + if err: + raise MyError() + return x + float("0.0") + +def test_indirect_call() -> None: + assert asyncio.run(indirect_call(3, inc)) == 4 + + with assertRaises(MyError): + asyncio.run(indirect_call_2(exc1())) + + assert asyncio.run(indirect_call_3(ident(2.0))) == 3.0 + assert asyncio.run(indirect_call_3(ident(-113.0))) == -112.0 + assert asyncio.run(indirect_call_3(ident(-114.0))) == -113.0 + + with assertRaises(MyError): + asyncio.run(indirect_call_3(ident(1.0, True))) + with assertRaises(MyError): + asyncio.run(indirect_call_3(ident(-113.0, True))) + [file asyncio/__init__.pyi] async def sleep(t: float) -> None: ... # eh, we could use the real type but it doesn't seem important From 3048d03bacb5a2e1522ef4e0be26aefff62bb662 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Thu, 3 Jul 2025 15:55:39 +0100 Subject: [PATCH 14/22] Support method calls --- mypyc/irbuild/statement.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 1d5be5385dd4c..2054d311fded2 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -932,7 +932,7 @@ def emit_yield_from_or_await( to_yield_reg = Register(object_rprimitive) received_reg = Register(object_rprimitive) - if isinstance(val, Call) and isinstance(val.type, RInstance) and val.type.class_ir.is_generated: + if isinstance(val, (Call, MethodCall)) and isinstance(val.type, RInstance) and val.type.class_ir.is_generated: # XXX is_generator not is_generated !!!!!!! FIXME FIXME iter_val = val else: From 2c184dcd6a53cbd82e03693fffae331697d4dfa4 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 4 Jul 2025 13:27:28 +0100 Subject: [PATCH 15/22] Add comments and docstrings --- mypyc/irbuild/mapper.py | 4 ++++ mypyc/irbuild/prepare.py | 6 ++++++ 2 files changed, 10 insertions(+) diff --git a/mypyc/irbuild/mapper.py b/mypyc/irbuild/mapper.py index 0fd59c6c7506c..99ae6db8c3bdd 100644 --- a/mypyc/irbuild/mapper.py +++ b/mypyc/irbuild/mapper.py @@ -64,6 +64,7 @@ def __init__(self, group_map: dict[str, str | None]) -> None: self.type_to_ir: dict[TypeInfo, ClassIR] = {} self.func_to_decl: dict[SymbolNode, FuncDecl] = {} self.symbol_fullnames: set[str] = set() + # The corresponding generator class that implements a generator/async function self.fdef_to_generator: dict[FuncDef, ClassIR] = {} def type_to_rtype(self, typ: Type | None) -> RType: @@ -174,6 +175,9 @@ def fdef_to_sig(self, fdef: FuncDef, strict_dunders_typing: bool) -> FuncSignatu arg_pos_onlys = [name is None for name in fdef.type.arg_names] # TODO: We could probably support decorators sometimes (static and class method?) if (fdef.is_coroutine or fdef.is_generator) and 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'. ret = RInstance(self.fdef_to_generator[fdef]) else: ret = self.type_to_rtype(fdef.type.ret_type) diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index 82be3dc91bdb6..3cd88ea357a55 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -196,6 +196,11 @@ def prepare_func_def( def create_generator_class_if_needed( module_name: str, class_name: str | None, fdef: FuncDef, mapper: Mapper, name_suffix: str = "" ) -> None: + """If function is a generator/async function, declare a generator class. + + Each generator and async function gets a dedicated class that implements the + generator protocol with generated methods. + """ if 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) @@ -213,6 +218,7 @@ def create_generator_class_if_needed( object_rprimitive, ) + # The implementation of most generator functionality is behind this magic method. helper_fn_decl = FuncDecl( "__mypyc_generator_helper__", name, From bcf9bbc575252dfa87c5be18feab883193df64e2 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 4 Jul 2025 13:27:58 +0100 Subject: [PATCH 16/22] Fix generator class detection --- mypyc/irbuild/statement.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 2054d311fded2..7a97308893846 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -932,8 +932,9 @@ def emit_yield_from_or_await( to_yield_reg = Register(object_rprimitive) received_reg = Register(object_rprimitive) - if isinstance(val, (Call, MethodCall)) and isinstance(val.type, RInstance) and val.type.class_ir.is_generated: - # XXX is_generator not is_generated !!!!!!! FIXME FIXME + helper_method = "__mypyc_generator_helper__" + if isinstance(val, (Call, MethodCall)) and isinstance(val.type, RInstance) and val.type.class_ir.has_method(helper_method): + # This is a generated generator class, and we can use a fast path. iter_val = val else: get_op = coro_op if is_await else iter_op @@ -946,7 +947,6 @@ def emit_yield_from_or_await( stop_block, main_block, done_block = BasicBlock(), BasicBlock(), BasicBlock() - helper_method = "__mypyc_generator_helper__" if isinstance(iter_reg.type, RInstance) and iter_reg.type.class_ir.has_method(helper_method): obj = builder.read(iter_reg) nn = builder.none_object() From ce74b2fd7edb210a78d6a6302e23682bebe20067 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 4 Jul 2025 13:28:49 +0100 Subject: [PATCH 17/22] Lint --- mypyc/irbuild/generator.py | 25 +++++++------------------ mypyc/irbuild/main.py | 12 +++++++++--- mypyc/irbuild/prepare.py | 18 ++++++++++-------- mypyc/irbuild/statement.py | 6 +++++- 4 files changed, 31 insertions(+), 30 deletions(-) diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 188c3982fdf25..c0a7bacf7347d 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -13,9 +13,9 @@ from typing import Callable from mypy.nodes import ARG_OPT, FuncDef, Var -from mypyc.common import ENV_ATTR_NAME, NEXT_LABEL_ATTR_NAME, SELF_NAME +from mypyc.common import ENV_ATTR_NAME, NEXT_LABEL_ATTR_NAME from mypyc.ir.class_ir import ClassIR -from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg +from mypyc.ir.func_ir import FuncDecl, FuncIR from mypyc.ir.ops import ( NO_TRACEBACK_LINE_NO, BasicBlock, @@ -78,9 +78,7 @@ def gen_generator_func( return func_ir, func_reg -def gen_generator_func_body( - builder: IRBuilder, fn_info: FuncInfo, func_reg: Value | None -) -> None: +def gen_generator_func_body(builder: IRBuilder, fn_info: FuncInfo, func_reg: Value | None) -> None: """Generate IR based on the body of a generator function. Add "__next__", "__iter__" and other generator methods to the generator @@ -229,10 +227,7 @@ def add_methods_to_generator_class( def add_helper_to_generator_class( - builder: IRBuilder, - arg_regs: list[Register], - blocks: list[BasicBlock], - fn_info: FuncInfo, + builder: IRBuilder, arg_regs: list[Register], blocks: list[BasicBlock], fn_info: FuncInfo ) -> FuncDecl: """Generates a helper method for a generator class, called by '__next__' and 'throw'.""" helper_fn_decl = fn_info.generator_class.ir.method_decls["__mypyc_generator_helper__"] @@ -252,9 +247,7 @@ def add_iter_to_generator_class(builder: IRBuilder, fn_info: FuncInfo) -> None: builder.add(Return(builder.self())) -def add_next_to_generator_class( - builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl -) -> None: +def add_next_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl) -> None: """Generates the '__next__' method for a generator class.""" with builder.enter_method(fn_info.generator_class.ir, "__next__", object_rprimitive, fn_info): none_reg = builder.none_object() @@ -269,9 +262,7 @@ def add_next_to_generator_class( builder.add(Return(result)) -def add_send_to_generator_class( - builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl -) -> None: +def add_send_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl) -> None: """Generates the 'send' method for a generator class.""" with builder.enter_method(fn_info.generator_class.ir, "send", object_rprimitive, fn_info): arg = builder.add_argument("arg", object_rprimitive) @@ -287,9 +278,7 @@ def add_send_to_generator_class( builder.add(Return(result)) -def add_throw_to_generator_class( - builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl -) -> None: +def add_throw_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl) -> None: """Generates the 'throw' method for a generator class.""" with builder.enter_method(fn_info.generator_class.ir, "throw", object_rprimitive, fn_info): typ = builder.add_argument("type", object_rprimitive) diff --git a/mypyc/irbuild/main.py b/mypyc/irbuild/main.py index eda5c0b0861b9..894e8f277723d 100644 --- a/mypyc/irbuild/main.py +++ b/mypyc/irbuild/main.py @@ -25,7 +25,7 @@ def f(x: int) -> int: from typing import Any, Callable, TypeVar, cast from mypy.build import Graph -from mypy.nodes import ClassDef, Expression, MypyFile, FuncDef +from mypy.nodes import ClassDef, Expression, FuncDef, MypyFile from mypy.state import state from mypy.types import Type from mypyc.analysis.attrdefined import analyze_always_defined_attrs @@ -37,7 +37,11 @@ def f(x: int) -> int: from mypyc.irbuild.builder import IRBuilder from mypyc.irbuild.mapper import Mapper from mypyc.irbuild.prebuildvisitor import PreBuildVisitor -from mypyc.irbuild.prepare import build_type_map, find_singledispatch_register_impls, create_generator_class_if_needed +from mypyc.irbuild.prepare import ( + build_type_map, + create_generator_class_if_needed, + find_singledispatch_register_impls, +) from mypyc.irbuild.visitor import IRBuilderVisitor from mypyc.irbuild.vtable import compute_vtable from mypyc.options import CompilerOptions @@ -81,7 +85,9 @@ def build_ir( if isinstance(fdef, FuncDef): # Make generator class name sufficiently unique. suffix = f"___{fdef.line}" - create_generator_class_if_needed(module.fullname, None, fdef, mapper, name_suffix=suffix) + create_generator_class_if_needed( + module.fullname, None, fdef, mapper, name_suffix=suffix + ) # Construct and configure builder objects (cyclic runtime dependency). visitor = IRBuilderVisitor() diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index 3cd88ea357a55..147392585b25e 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -38,7 +38,7 @@ from mypy.semanal import refers_to_fullname from mypy.traverser import TraverserVisitor from mypy.types import Instance, Type, get_proper_type -from mypyc.common import PROPSET_PREFIX, get_id_from_name, SELF_NAME +from mypyc.common import PROPSET_PREFIX, SELF_NAME, get_id_from_name from mypyc.crash import catch_errors from mypyc.errors import Errors from mypyc.ir.class_ir import ClassIR @@ -51,7 +51,14 @@ RuntimeArg, ) from mypyc.ir.ops import DeserMaps -from mypyc.ir.rtypes import RInstance, RType, dict_rprimitive, none_rprimitive, tuple_rprimitive +from mypyc.ir.rtypes import ( + RInstance, + RType, + dict_rprimitive, + none_rprimitive, + object_rprimitive, + tuple_rprimitive, +) from mypyc.irbuild.mapper import Mapper from mypyc.irbuild.util import ( get_func_def, @@ -62,7 +69,6 @@ ) from mypyc.options import CompilerOptions from mypyc.sametype import is_same_type -from mypyc.ir.rtypes import object_rprimitive def build_type_map( @@ -220,11 +226,7 @@ def create_generator_class_if_needed( # The implementation of most generator functionality is behind this magic method. helper_fn_decl = FuncDecl( - "__mypyc_generator_helper__", - name, - module_name, - helper_sig, - internal=True, + "__mypyc_generator_helper__", name, module_name, helper_sig, internal=True ) cir.method_decls[helper_fn_decl.name] = helper_fn_decl diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 7a97308893846..4bce8ecb210fc 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -933,7 +933,11 @@ def emit_yield_from_or_await( received_reg = Register(object_rprimitive) helper_method = "__mypyc_generator_helper__" - if isinstance(val, (Call, MethodCall)) and isinstance(val.type, RInstance) and val.type.class_ir.has_method(helper_method): + if ( + isinstance(val, (Call, MethodCall)) + and isinstance(val.type, RInstance) + and val.type.class_ir.has_method(helper_method) + ): # This is a generated generator class, and we can use a fast path. iter_val = val else: From 0ab4b19c11089f1ef3b6b2f7809922935974662a Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 4 Jul 2025 14:01:41 +0100 Subject: [PATCH 18/22] Test calling async method --- mypyc/test-data/run-async.test | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test index 99276688465f3..b8c4c22daf719 100644 --- a/mypyc/test-data/run-async.test +++ b/mypyc/test-data/run-async.test @@ -107,6 +107,29 @@ def test_indirect_call() -> None: with assertRaises(MyError): asyncio.run(indirect_call_3(ident(-113.0, True))) +class C: + def __init__(self, n: int) -> None: + self.n = n + + async def add(self, x: int, err: bool = False) -> int: + await asyncio.sleep(0) + if err: + raise MyError() + return x + self.n + +async def method_call(x: int) -> int: + c = C(5) + return await c.add(x) + +async def method_call_exception() -> int: + c = C(5) + return await c.add(3, err=True) + +def test_async_method_call() -> None: + assert asyncio.run(method_call(3)) == 8 + with assertRaises(MyError): + asyncio.run(method_call_exception()) + [file asyncio/__init__.pyi] async def sleep(t: float) -> None: ... # eh, we could use the real type but it doesn't seem important From 11a1e7513ff54539fbfc7bb879b9e9e98f671c1b Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 4 Jul 2025 14:05:34 +0100 Subject: [PATCH 19/22] Add comment --- mypyc/irbuild/statement.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 4bce8ecb210fc..46b19a1b3b9e1 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -955,6 +955,7 @@ def emit_yield_from_or_await( obj = builder.read(iter_reg) nn = builder.none_object() m = MethodCall(obj, helper_method, [nn, nn, nn, nn], line) + # Generators have custom error handling, so disable normal error handling. m.error_kind = ERR_NEVER _y_init = builder.add(m) else: From fff29dfafce708c3285d5c38cf6a1577ef5fa860 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 4 Jul 2025 14:07:33 +0100 Subject: [PATCH 20/22] Add some type annotations --- mypyc/irbuild/mapper.py | 2 +- mypyc/irbuild/statement.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/mypyc/irbuild/mapper.py b/mypyc/irbuild/mapper.py index 99ae6db8c3bdd..4a01255e2d5d5 100644 --- a/mypyc/irbuild/mapper.py +++ b/mypyc/irbuild/mapper.py @@ -178,7 +178,7 @@ def fdef_to_sig(self, fdef: FuncDef, strict_dunders_typing: bool) -> FuncSignatu # 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'. - ret = RInstance(self.fdef_to_generator[fdef]) + ret: RType = RInstance(self.fdef_to_generator[fdef]) else: ret = self.type_to_rtype(fdef.type.ret_type) else: diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 46b19a1b3b9e1..982adc6928b08 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -939,7 +939,7 @@ def emit_yield_from_or_await( and val.type.class_ir.has_method(helper_method) ): # This is a generated generator class, and we can use a fast path. - iter_val = val + iter_val: Value = val else: get_op = coro_op if is_await else iter_op if isinstance(get_op, PrimitiveDescription): From c33a683132881839b86b9d4cc2159dbee4a6be40 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 4 Jul 2025 14:31:25 +0100 Subject: [PATCH 21/22] Fix self check --- mypyc/irbuild/generator.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index c0a7bacf7347d..eec27e1cfb84d 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -152,6 +152,7 @@ def instantiate_generator_class(builder: IRBuilder) -> Value: def setup_generator_class(builder: IRBuilder) -> ClassIR: mapper = builder.mapper + assert isinstance(builder.fn_info.fitem, FuncDef) generator_class_ir = mapper.fdef_to_generator[builder.fn_info.fitem] if builder.fn_info.can_merge_generator_and_env_classes(): builder.fn_info.env_class = generator_class_ir From 1fa32b2557da78de8d95761978f2bd5b3df2ee4a Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Mon, 7 Jul 2025 14:41:37 +0100 Subject: [PATCH 22/22] Add comments based on review feedback --- mypyc/irbuild/statement.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 982adc6928b08..9c7ffb6a3adf9 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -938,7 +938,12 @@ def emit_yield_from_or_await( and isinstance(val.type, RInstance) and val.type.class_ir.has_method(helper_method) ): - # This is a generated generator class, and we can use a fast path. + # This is a generated native generator class, and we can use a fast path. + # This allows two optimizations: + # 1) No need to call CPy_GetCoro() or iter() since for native generators + # it just returns the generator object (implemented here). + # 2) Instead of calling next(), call generator helper method directly, + # since next() just calls __next__ which calls the helper method. iter_val: Value = val else: get_op = coro_op if is_await else iter_op @@ -952,6 +957,7 @@ def emit_yield_from_or_await( stop_block, main_block, done_block = BasicBlock(), BasicBlock(), BasicBlock() if isinstance(iter_reg.type, RInstance) and iter_reg.type.class_ir.has_method(helper_method): + # Second fast path optimization: call helper directly (see also comment above). obj = builder.read(iter_reg) nn = builder.none_object() m = MethodCall(obj, helper_method, [nn, nn, nn, nn], line)