Make it impossible to extract if return is not at the end

jedi/api/refactoring/extract.py

@@ -236,7 +236,7 @@ def extract_function(inference_state, path, module_context, name, pos, until_pos
                 nodes[0].parent,
                 nodes[-1].end_pos,
                 return_variables
-            )) or [return_variables[-1]]
+            )) or [return_variables[-1]] if return_variables else []
 
         remaining_prefix, code_block = _suite_nodes_to_string(nodes, pos)
         after_leaf = nodes[-1].get_next_leaf()
@@ -297,7 +297,11 @@ def _check_for_non_extractables(nodes):
         try:
             children = n.children
         except AttributeError:
-            pass
+            if n.value in ('return', 'yield'):
+                raise RefactoringError(
+                    'Can only extract %s statements if they are at the end.'
+                    % n.value
+                )
         else:
             _check_for_non_extractables(children)
 
@@ -306,10 +310,8 @@ def _is_name_input(module_context, names, first, last):
     for name in names:
         if name.api_type == 'param' or not name.parent_context.is_module():
             if name.get_root_context() is not module_context:
-                print('true')
                 return True
             if name.start_pos is None or not (first <= name.start_pos < last):
-                print('true1', first, name.start_pos, last)
                 return True
     return False
 

test/refactor/extract_function.py

@@ -395,3 +395,38 @@ class X:
         return self.ab(local1, b)
         # bar
     local2
+# -------------------------------------------------- random-return-1
+def x():
+    #? 0 error {'new_name': 'ab', 'until_line': 5, 'until_column': 10}
+    if x:
+        return 1
+    return 1
+# ++++++++++++++++++++++++++++++++++++++++++++++++++
+Can only extract return statements if they are at the end.
+# -------------------------------------------------- random-return-2
+def x():
+    #? 0 error {'new_name': 'ab', 'until_line': 5, 'until_column': 10}
+    #
+    return
+    pass
+# ++++++++++++++++++++++++++++++++++++++++++++++++++
+Can only extract return statements if they are at the end.
+# -------------------------------------------------- random-yield-1
+def x():
+    #? 0 error {'new_name': 'ab', 'until_line': 5, 'until_column': 10}
+    #
+    if (yield 1):
+        return
+    pass
+# ++++++++++++++++++++++++++++++++++++++++++++++++++
+Can only extract yield statements if they are at the end.
+# -------------------------------------------------- random-yield-2
+def x():
+    #? 0 error {'new_name': 'ab', 'until_line': 4, 'until_column': 10}
+    #
+    try:
+        yield
+    finally:
+        pass
+# ++++++++++++++++++++++++++++++++++++++++++++++++++
+Can only extract yield statements if they are at the end.