Refactor

python · JukkaL · May 2, 2025 · Apr 30, 2025 · Apr 30, 2025 · Apr 30, 2025
commit 476ebe1e858dc9f20dd5949ef251456d57776fa1
diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py
@@ -1169,7 +1169,7 @@ def flatten_classes(self, arg: RefExpr | TupleExpr) -> list[ClassIR] | None:
                     return None
             return res
 
-    def enter(self, fn_info: FuncInfo | str = "") -> None:
+    def enter(self, fn_info: FuncInfo | str = "", *, ret_type: RType = none_rprimitive) -> None:
         if isinstance(fn_info, str):
             fn_info = FuncInfo(name=fn_info)
         self.builder = LowLevelIRBuilder(self.errors, self.options)
@@ -1179,7 +1179,7 @@ def enter(self, fn_info: FuncInfo | str = "") -> None:
         self.runtime_args.append([])
         self.fn_info = fn_info
         self.fn_infos.append(self.fn_info)
-        self.ret_types.append(none_rprimitive)
+        self.ret_types.append(ret_type)
         if fn_info.is_generator:
             self.nonlocal_control.append(GeneratorNonlocalControl())
         else:

diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
@@ -235,20 +235,20 @@ def c() -> None:
         func_name = singledispatch_main_func_name(name)
     else:
         func_name = name
-    builder.enter(
-        FuncInfo(
-            fitem=fitem,
-            name=func_name,
-            class_name=class_name,
-            namespace=gen_func_ns(builder),
-            is_nested=is_nested,
-            contains_nested=contains_nested,
-            is_decorated=is_decorated,
-            in_non_ext=in_non_ext,
-            add_nested_funcs_to_env=add_nested_funcs_to_env,
-        )
+
+    fn_info = FuncInfo(
+        fitem=fitem,
+        name=func_name,
+        class_name=class_name,
+        namespace=gen_func_ns(builder),
+        is_nested=is_nested,
+        contains_nested=contains_nested,
+        is_decorated=is_decorated,
+        in_non_ext=in_non_ext,
+        add_nested_funcs_to_env=add_nested_funcs_to_env,
     )
-    is_generator = builder.fn_info.is_generator
+    is_generator = fn_info.is_generator
+    builder.enter(fn_info, ret_type=sig.ret_type if not is_generator else object_rprimitive)
 
     # 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
@@ -266,6 +266,10 @@ def c() -> None:
             builder,
             lambda args, blocks, fn_info: gen_func_ir(builder, args, blocks, sig, fn_info, cdef, is_singledispatch))
 
+        # Re-enter the FuncItem and visit the body of the function this time.
+        builder.enter(fn_info, ret_type=sig.ret_type)
+        setup_env_for_generator_class(builder)
+
         load_outer_envs(builder, builder.fn_info.generator_class)
         top_level = builder.top_level_fn_info()
         if (
@@ -283,8 +287,6 @@ def c() -> None:
         if contains_nested:
             finalize_env_class(builder)
 
-    builder.ret_types[-1] = sig.ret_type
-
     add_vars_to_env(builder)
 
     builder.accept(fitem.body)

diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py
@@ -64,10 +64,6 @@ def gen_generator_func(
     args, _, blocks, ret_type, fn_info = builder.leave()
     func_ir, func_reg = gen_func_ir(args, blocks, fn_info)
 
-    # Re-enter the FuncItem and visit the body of the function this time.
-    builder.enter(fn_info)
-    setup_env_for_generator_class(builder)
-
     return func_ir, func_reg