Looped if then on 3.6-3.8, Looped Return on all versions

pylingual/control_flow_reconstruction/templates/Conditional.py

@@ -1,5 +1,5 @@
 from ..cft import ControlFlowTemplate, EdgeKind, register_template
-from ..utils import T, N, defer_source_to, run_is, has_no_lines, with_instructions, has_instval, starting_instructions, to_indented_source, make_try_match, without_top_level_instructions
+from ..utils import T, N, defer_source_to, run_is, has_no_lines, with_instructions, exact_instructions, without_instructions, has_instval, starting_instructions, to_indented_source, make_try_match, without_top_level_instructions
 from .Loop import BreakTemplate, ContinueTemplate
 
 
@@ -51,7 +51,7 @@ class IfElseLoop(ControlFlowTemplate):
 class IfThen(ControlFlowTemplate):
     template = T(
         if_header=~N("if_body", "tail").with_cond(without_top_level_instructions("WITH_EXCEPT_START", "CHECK_EXC_MATCH", "FOR_ITER")),
-        if_body=~N("tail").with_in_deg(1) | ~N("tail.").with_in_deg(1).with_cond(run_is(2)),
+        if_body=~N("tail").with_in_deg(1) | ~N("tail.").with_in_deg(1).with_cond(run_is(2)) | ~N.tail().with_in_deg(1).with_cond(exact_instructions("LOAD_CONST","RETURN_VALUE")),
         tail=N.tail(),
     )
 

pylingual/control_flow_reconstruction/templates/Loop.py

@@ -13,6 +13,7 @@ from ..utils import (
     versions_from,
     with_instructions,
     exact_instructions,
+    ending_instructions,
     has_no_lines,
     has_some_lines,
     condense_mapping,
@@ -45,6 +46,22 @@ class ForLoop(ControlFlowTemplate):
             {for_body}
         """
 
+@register_template(0, 2)
+class LoopedReturn(ControlFlowTemplate):
+    template = T(
+        for_iter=~N("for_body", "tail").with_cond(ending_instructions("FOR_ITER")),
+        for_body=~N.tail().with_in_deg(1).with_cond(ending_instructions("RETURN_CONST"),ending_instructions("RETURN_VALUE")),
+        tail=N.tail(),
+    )
+
+    try_match = make_try_match({EdgeKind.Fall: "tail"}, "for_iter", "for_body")
+
+    @to_indented_source
+    def to_indented_source():
+        """
+        {for_iter}
+            {for_body}
+        """
 
 @register_template(0, 2, *versions_below(3, 10))
 class SelfLoop3_6(ControlFlowTemplate):