Polish

Author: JukkaL

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

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)