From e22ec9c7ecd1960185fbbfe2c4081ed6b12d35bc Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Thu, 9 Oct 2025 13:51:27 +0100 Subject: [PATCH 1/3] WIP refactor generator class creation to a separate pass --- mypyc/irbuild/main.py | 2 ++ mypyc/irbuild/prepare.py | 9 ++++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/main.py b/mypyc/irbuild/main.py index d2c8924a7298..ff9b535b0e2f 100644 --- a/mypyc/irbuild/main.py +++ b/mypyc/irbuild/main.py @@ -39,6 +39,7 @@ def f(x: int) -> int: from mypyc.irbuild.prebuildvisitor import PreBuildVisitor from mypyc.irbuild.prepare import ( build_type_map, + gen_generator_types, create_generator_class_if_needed, find_singledispatch_register_impls, ) @@ -68,6 +69,7 @@ def build_ir( """ build_type_map(mapper, modules, graph, types, options, errors) + gen_generator_types(mapper, modules) singledispatch_info = find_singledispatch_register_impls(modules, errors) result: ModuleIRs = {} diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index e4f43b38b0dc..37406b00a379 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -196,7 +196,8 @@ def prepare_func_def( mapper: Mapper, options: CompilerOptions, ) -> FuncDecl: - create_generator_class_if_needed(module_name, class_name, fdef, mapper) + #create_generator_class_if_needed(module_name, class_name, fdef, mapper) + #assert False kind = ( FUNC_CLASSMETHOD @@ -811,3 +812,9 @@ def registered_impl_from_possible_register_call( if isinstance(node, Decorator): return RegisteredImpl(node.func, dispatch_type) return None + + +def gen_generator_types(mapper: Mapper, modules: list[MypyFile]) -> None: + for fdef, ir in mapper.func_to_decl.items(): + if isinstance(fdef, FuncDef): + create_generator_class_if_needed(ir.module_name, ir.class_name, fdef, mapper) From 3fd3b1ce840193a05090b61e73abf4066ac3f194 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Thu, 9 Oct 2025 14:16:04 +0100 Subject: [PATCH 2/3] Fix --- mypyc/irbuild/main.py | 12 ++++---- mypyc/irbuild/mapper.py | 9 +----- mypyc/irbuild/prepare.py | 65 ++++++++++++++++++++++------------------ 3 files changed, 44 insertions(+), 42 deletions(-) diff --git a/mypyc/irbuild/main.py b/mypyc/irbuild/main.py index ff9b535b0e2f..02ef5c5dac54 100644 --- a/mypyc/irbuild/main.py +++ b/mypyc/irbuild/main.py @@ -39,9 +39,9 @@ def f(x: int) -> int: from mypyc.irbuild.prebuildvisitor import PreBuildVisitor from mypyc.irbuild.prepare import ( build_type_map, - gen_generator_types, - create_generator_class_if_needed, + create_generator_class_for_func, find_singledispatch_register_impls, + gen_generator_types, ) from mypyc.irbuild.visitor import IRBuilderVisitor from mypyc.irbuild.vtable import compute_vtable @@ -89,9 +89,11 @@ 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 - ) + # TODO: decorated? + if fdef.is_coroutine or fdef.is_generator: + create_generator_class_for_func( + module.fullname, None, fdef, mapper, name_suffix=suffix + ) # Construct and configure builder objects (cyclic runtime dependency). visitor = IRBuilderVisitor() diff --git a/mypyc/irbuild/mapper.py b/mypyc/irbuild/mapper.py index 05aa0e45c569..c986499b6f65 100644 --- a/mypyc/irbuild/mapper.py +++ b/mypyc/irbuild/mapper.py @@ -180,14 +180,7 @@ 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] - # 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: RType = RInstance(self.fdef_to_generator[fdef]) - else: - ret = self.type_to_rtype(fdef.type.ret_type) + 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 37406b00a379..e9b92cf4369d 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -196,8 +196,8 @@ def prepare_func_def( mapper: Mapper, options: CompilerOptions, ) -> FuncDecl: - #create_generator_class_if_needed(module_name, class_name, fdef, mapper) - #assert False + # create_generator_class_if_needed(module_name, class_name, fdef, mapper) + # assert False kind = ( FUNC_CLASSMETHOD @@ -210,38 +210,37 @@ def prepare_func_def( return decl -def create_generator_class_if_needed( +def create_generator_class_for_func( 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. +) -> ClassIR: + """For 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) - cir.reuse_freed_instance = True - mapper.fdef_to_generator[fdef] = cir + 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.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), - # If non-NULL, used to store return value instead of raising StopIteration(retv) - RuntimeArg("stop_iter_ptr", object_pointer_rprimitive), - ), - object_rprimitive, - ) + helper_sig = FuncSignature( + ( + RuntimeArg(SELF_NAME, object_rprimitive), + RuntimeArg("type", object_rprimitive), + RuntimeArg("value", object_rprimitive), + RuntimeArg("traceback", object_rprimitive), + RuntimeArg("arg", object_rprimitive), + # If non-NULL, used to store return value instead of raising StopIteration(retv) + RuntimeArg("stop_iter_ptr", object_pointer_rprimitive), + ), + object_rprimitive, + ) - # The implementation of most generator functionality is behind this magic method. - helper_fn_decl = FuncDecl( - GENERATOR_HELPER_NAME, name, module_name, helper_sig, internal=True - ) - cir.method_decls[helper_fn_decl.name] = helper_fn_decl + # The implementation of most generator functionality is behind this magic method. + helper_fn_decl = FuncDecl(GENERATOR_HELPER_NAME, name, module_name, helper_sig, internal=True) + cir.method_decls[helper_fn_decl.name] = helper_fn_decl + return cir def prepare_method_def( @@ -816,5 +815,13 @@ def registered_impl_from_possible_register_call( def gen_generator_types(mapper: Mapper, modules: list[MypyFile]) -> None: for fdef, ir in mapper.func_to_decl.items(): - if isinstance(fdef, FuncDef): - create_generator_class_if_needed(ir.module_name, ir.class_name, fdef, mapper) + 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) + # 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) From a0e4f431e32fd3921f567001cc48e1bf268f3641 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 10 Oct 2025 12:51:32 +0100 Subject: [PATCH 3/3] Polish --- mypyc/irbuild/main.py | 5 ++--- mypyc/irbuild/prepare.py | 10 ++++++---- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/mypyc/irbuild/main.py b/mypyc/irbuild/main.py index 02ef5c5dac54..f08911a1bc4c 100644 --- a/mypyc/irbuild/main.py +++ b/mypyc/irbuild/main.py @@ -38,10 +38,10 @@ def f(x: int) -> int: from mypyc.irbuild.mapper import Mapper from mypyc.irbuild.prebuildvisitor import PreBuildVisitor from mypyc.irbuild.prepare import ( + adjust_generator_classes_of_methods, build_type_map, create_generator_class_for_func, find_singledispatch_register_impls, - gen_generator_types, ) from mypyc.irbuild.visitor import IRBuilderVisitor from mypyc.irbuild.vtable import compute_vtable @@ -69,7 +69,7 @@ def build_ir( """ build_type_map(mapper, modules, graph, types, options, errors) - gen_generator_types(mapper, modules) + adjust_generator_classes_of_methods(mapper) singledispatch_info = find_singledispatch_register_impls(modules, errors) result: ModuleIRs = {} @@ -89,7 +89,6 @@ def build_ir( if isinstance(fdef, FuncDef): # Make generator class name sufficiently unique. suffix = f"___{fdef.line}" - # TODO: decorated? if fdef.is_coroutine or fdef.is_generator: create_generator_class_for_func( module.fullname, None, fdef, mapper, name_suffix=suffix diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index e9b92cf4369d..2d0a1a8f03bf 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -196,9 +196,6 @@ def prepare_func_def( mapper: Mapper, options: CompilerOptions, ) -> FuncDecl: - # create_generator_class_if_needed(module_name, class_name, fdef, mapper) - # assert False - kind = ( FUNC_CLASSMETHOD if fdef.is_class @@ -813,7 +810,12 @@ def registered_impl_from_possible_register_call( return None -def gen_generator_types(mapper: Mapper, modules: list[MypyFile]) -> None: +def adjust_generator_classes_of_methods(mapper: Mapper) -> None: + """Make optimizations and adjustments to generated generator classes of methods. + + 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(): 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)