Make sure that return at the end works properly for extract

jedi/api/refactoring/extract.py

@@ -224,6 +224,7 @@ def extract_function(inference_state, path, module_context, name, pos, until_pos
     if is_expression:
         code_block = 'return ' + _expression_nodes_to_string(nodes) + '\n'
         remaining_prefix = None
+        has_ending_return_stmt = False
     else:
         has_ending_return_stmt = _is_node_ending_return_stmt(nodes[-1])
         if not has_ending_return_stmt:
@@ -247,6 +248,9 @@ def extract_function(inference_state, path, module_context, name, pos, until_pos
             output_var_str = ', '.join(return_variables)
             code_block += 'return ' + output_var_str + '\n'
 
+    # Check if we have to raise RefactoringError
+    _check_for_non_extractables(nodes[:-1] if has_ending_return_stmt else nodes)
+
     decorator = ''
     self_param = None
     if is_bound_method:
@@ -288,6 +292,16 @@ def extract_function(inference_state, path, module_context, name, pos, until_pos
     return Refactoring(inference_state.grammar, file_to_node_changes)
 
 
+def _check_for_non_extractables(nodes):
+    for n in nodes:
+        try:
+            children = n.children
+        except AttributeError:
+            pass
+        else:
+            _check_for_non_extractables(children)
+
+
 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():

test/refactor/extract_function.py

@@ -367,7 +367,31 @@ class X:
         local1 = 3
         local2 = 4
         return local1 * glob1 * b
+        # bar
 
     def f(self, b, c):
         #? 11 text {'new_name': 'ab', 'until_line': 11, 'until_column': 10}
         return self.ab(b)
+# -------------------------------------------------- in-method-range-3
+glob1 = 1
+class X:
+    def f(self, b, c):
+        local1, local2 = 3, 4
+        #foo
+        #? 11 text {'new_name': 'ab', 'until_line': 7, 'until_column': 29}
+        return local1 & glob1 & b
+        # bar
+    local2
+# ++++++++++++++++++++++++++++++++++++++++++++++++++
+glob1 = 1
+class X:
+    def ab(self, local1, b):
+        return local1 & glob1 & b
+
+    def f(self, b, c):
+        local1, local2 = 3, 4
+        #foo
+        #? 11 text {'new_name': 'ab', 'until_line': 7, 'until_column': 29}
+        return self.ab(local1, b)
+        # bar
+    local2