Make sure that extract variable works for some ranges

jedi/api/__init__.py

@@ -562,12 +562,13 @@ class Script(object):
         return self._extract_variable(line, column, **kwargs)  # Python 2...
 
     def _extract_variable(self, line, column, new_name, until_line=None, until_column=None):
-        if until_line is None or until_column is None:
-            if until_line is not until_column:
-                raise TypeError('If you provide until_line or until_column '
-                                'you have to provide both')
+        if until_line is None and until_column is None:
             until_pos = None
         else:
+            if until_line is None:
+                until_line = line
+            if until_line is None:
+                raise TypeError('If you provide until_line you have to provide until_column')
             until_pos = until_line, until_column
         return refactoring.extract_variable(
             self._grammar, self.path, self._module_node, new_name, (line, column), until_pos)

jedi/api/refactoring.py

@@ -232,7 +232,6 @@ def extract_variable(grammar, path, module_node, new_name, pos, until_pos):
         while node.parent.type in _EXTRACT_USE_PARENT:
             node = node.parent
         start_leaf
-        extracted = node.get_code(include_prefix=False)
         nodes = [node]
     else:
         end_leaf = module_node.get_leaf_for_position(until_pos, include_prefixes=True)
@@ -240,16 +239,21 @@ def extract_variable(grammar, path, module_node, new_name, pos, until_pos):
             end_leaf = end_leaf.get_previous_leaf()
             if end_leaf is None:
                 raise RefactoringError('Cannot extract anything from that')
+        parent_node = start_leaf
+        while parent_node.end_pos < end_leaf.end_pos:
+            parent_node = parent_node.parent
+        nodes = [parent_node]
     if any(node.type == 'name' and node.is_definition() for node in nodes):
         raise RefactoringError('Cannot extract a definition of a name')
     if any(node.type in _NON_EXCTRACABLE for node in nodes) \
             or nodes[0].type == 'keyword' and nodes[0].value not in ('None', 'True', 'False'):
         raise RefactoringError('Cannot extract a %s' % node.type)
 
-    definition = _get_parent_definition(node)
+    definition = _get_parent_definition(nodes[0])
     first_definition_leaf = definition.get_first_leaf()
 
     dct = {}
+    extracted = ''.join(n.get_code(include_prefix=i != 0) for i, n in enumerate(nodes))
     extracted_prefix = _insert_line_before(
         first_definition_leaf.prefix,
         new_name + ' = ' + extracted,

test/refactor/extract_variable.py

@@ -119,3 +119,38 @@ x +  1, 3
 #? 4 text {'new_name': 'x'}
 x = x +  1
 x, 3
+# -------------------------------------------------- range-1
+#? 4 text {'new_name': 'x', 'until_column': 9}
+y +  1, 3
+# ++++++++++++++++++++++++++++++++++++++++++++++++++
+#? 4 text {'new_name': 'x', 'until_column': 9}
+x = y +  1, 3
+x
+# -------------------------------------------------- range-2
+#? 1 text {'new_name': 'x', 'until_column': 3}
+y +  1, 3
+# ++++++++++++++++++++++++++++++++++++++++++++++++++
+#? 1 text {'new_name': 'x', 'until_column': 3}
+x = y +  1
+x, 3
+# -------------------------------------------------- range-3
+#? 1 text {'new_name': 'x', 'until_column': 6}
+y +  1, 3
+# ++++++++++++++++++++++++++++++++++++++++++++++++++
+#? 1 text {'new_name': 'x', 'until_column': 6}
+x = y +  1
+x, 3
+# -------------------------------------------------- range-4
+#? 1 text {'new_name': 'x', 'until_column': 1}
+y +  1, 3
+# ++++++++++++++++++++++++++++++++++++++++++++++++++
+#? 1 text {'new_name': 'x', 'until_column': 1}
+x = y
+x +  1, 3
+# -------------------------------------------------- range-5
+#? 0 text {'new_name': 'x', 'until_column': 1}
+y +  1, 3
+# ++++++++++++++++++++++++++++++++++++++++++++++++++
+#? 0 text {'new_name': 'x', 'until_column': 1}
+x = y
+x +  1, 3