Fix: Extract can now deal with return statements at the end

jedi/api/refactoring/extract.py

@@ -207,6 +207,8 @@ def _is_not_extractable_syntax(node):
 
 def extract_function(inference_state, path, module_context, name, pos, until_pos):
     nodes = _find_nodes(module_context.tree_node, pos, until_pos)
+    assert len(nodes)
+
     is_expression, _ = _is_expression_with_error(nodes)
     context = module_context.create_context(nodes[0])
     is_bound_method = context.is_bound_method()
@@ -223,6 +225,8 @@ def extract_function(inference_state, path, module_context, name, pos, until_pos
         code_block = 'return ' + _expression_nodes_to_string(nodes) + '\n'
         remaining_prefix = None
     else:
+        has_ending_return_stmt = _is_node_ending_return_stmt(nodes[-1])
+        if not has_ending_return_stmt:
             # Find the actually used variables (of the defined ones). If none are
             # used (e.g. if the range covers the whole function), return the last
             # defined variable.
@@ -239,6 +243,7 @@ def extract_function(inference_state, path, module_context, name, pos, until_pos
         code_block += first
 
         code_block = dedent(code_block)
+        if not has_ending_return_stmt:
             output_var_str = ', '.join(return_variables)
             code_block += 'return ' + output_var_str + '\n'
 
@@ -269,6 +274,9 @@ def extract_function(inference_state, path, module_context, name, pos, until_pos
     )
     if is_expression:
         replacement = function_call
+    else:
+        if has_ending_return_stmt:
+            replacement = 'return ' + function_call + '\n'
         else:
             replacement = output_var_str + ' = ' + function_call + '\n'
 
@@ -353,3 +361,10 @@ def _find_needed_output_variables(context, search_node, at_least_pos, return_var
             if not name.is_definition() and name.value in return_variables:
                 return_variables.remove(name.value)
                 yield name.value
+
+
+def _is_node_ending_return_stmt(node):
+    t = node.type
+    if t == 'simple_stmt':
+        return _is_node_ending_return_stmt(node.children[0])
+    return t == 'return_stmt'

test/refactor/extract_function.py

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