Fixed all on_import tests.

jedi/api/completion.py

@@ -60,6 +60,7 @@ class Completion:
         self._parser = parser
         self._module = evaluator.wrap(parser.module())
         self._user_context = user_context
+        self._source = user_context.source
         self._pos = position
         self._call_signatures_method = call_signatures_method
 
@@ -132,7 +133,7 @@ class Completion:
         if completion_parts.name:
             pos = pos[0], pos[1] - len(completion_parts.name)
 
-        stack = helpers.get_stack_at_position(grammar, self._module, pos)
+        stack = helpers.get_stack_at_position(grammar, self._source, self._module, pos)
         allowed_keywords, allowed_tokens = \
             helpers.get_possible_completion_types(grammar, stack)
 

jedi/api/helpers.py

@@ -50,55 +50,61 @@ def check_error_statements(module, pos):
     return None, 0, False, False
 
 
-def get_code_until(code, code_start_pos, end_pos):
+def get_code(code, start_pos, end_pos):
     """
     :param code_start_pos: is where the code starts.
     """
     lines = common.splitlines(code)
-    line_difference = end_pos[0] - code_start_pos[0]
+    # Get relevant lines.
-    if line_difference == 0:
+    lines = lines[start_pos[0] - 1:end_pos[0]]
-        end_line_length = end_pos[1] - code_start_pos[1]
+    # Remove the parts at the end of the line.
-    else:
+    lines[-1] = lines[-1][:end_pos[1]]
-        end_line_length = end_pos[1]
+    # Remove first line indentation.
-
+    lines[0] = lines[0][start_pos[1]:]
-    if line_difference > len(lines) or end_line_length > len(lines[line_difference]):
+    return ''.join(lines)
-        raise ValueError("The end_pos seems to be after the code part.")
-
-    new_lines = lines[:line_difference] + [lines[line_difference][:end_line_length]]
-    return '\n'.join(new_lines)
 
 
-def get_stack_at_position(grammar, module, pos):
+def get_user_or_error_stmt(module, position):
-    """
+    user_stmt = module.get_statement_for_position(position)
-    Returns the possible node names (e.g. import_from, xor_test or yield_stmt).
-    """
-    user_stmt = module.get_statement_for_position(pos)
     if user_stmt is None or user_stmt.type == 'whitespace':
         # If there's no error statement and we're just somewhere, we want
         # completions for just whitespace.
-        code = ''
+        for error_stmt in module.error_statements:
+            if error_stmt.start_pos < position <= error_stmt.end_pos:
+                return error_stmt
 
-        for error_statement in module.error_statements:
+    return user_stmt
-            if error_statement.start_pos < pos <= error_statement.end_pos:
+
-                code = error_statement.get_code(include_prefix=False)
+
-                node = error_statement
+def get_stack_at_position(grammar, source, module, pos):
+    """
+    Returns the possible node names (e.g. import_from, xor_test or yield_stmt).
+    """
+    user_stmt = get_user_or_error_stmt(module, pos)
+
+    if user_stmt is None:
+        user_stmt = module.get_leaf_for_position(pos, include_prefixes=True)
+    # Only if were in front of the leaf we want to get the stack,
+    # because after there's probably a newline or whatever that would
+    # be actually tokenized and is not just prefix.
+    if pos <= user_stmt.start_pos:
+        leaf = user_stmt.get_previous_leaf()
+        for error_stmt in reversed(module.error_statements):
+            if leaf.start_pos <= error_stmt.start_pos <= user_stmt.start_pos:
+                # The leaf appears not to be the last leaf. It's actually an
+                # error statement.
+                user_stmt = error_stmt
                 break
         else:
-            raise NotImplementedError
+            user_stmt = get_user_or_error_stmt(module, leaf.start_pos)
-    else:
-        code = user_stmt.get_code_with_error_statements(include_prefix=False)
-        node = user_stmt
 
-    # Make sure we include the whitespace after the statement as well, since it
-    # could be where we would want to complete.
-    print('a', repr(code), node, repr(node.get_next_leaf().prefix))
-    code += node.get_next_leaf().prefix
 
-    code = get_code_until(code, node.start_pos, pos)
+    print(user_stmt.start_pos, pos)
+    code = get_code(source, user_stmt.start_pos, pos)
     # Remove whitespace at the end. Necessary, because the tokenizer will parse
     # an error token (there's no new line at the end in our case). This doesn't
     # alter any truth about the valid tokens at that position.
-    code = code.rstrip()
+    code = code.strip()
 
     class EndMarkerReached(Exception):
         pass
@@ -110,9 +116,8 @@ def get_stack_at_position(grammar, module, pos):
             else:
                 yield token_
 
-    print(repr(code))
+    print(repr(code), 'x')
-    p = parser.Parser(grammar, code,
+    p = parser.Parser(grammar, code, start_parsing=False)
-                      start_parsing=False)
     try:
         p.parse(tokenizer=tokenize_without_endmarker(code))
     except EndMarkerReached:
@@ -147,17 +152,25 @@ def get_possible_completion_types(grammar, stack):
                 for first_label_index in itsfirst.keys():
                     add_results(first_label_index)
 
-    dfa, state, node = stack[-1]
-    states, first = dfa
-    arcs = states[state]
-
     inversed_keywords = dict((v, k) for k, v in grammar.keywords.items())
     inversed_tokens = dict((v, k) for k, v in grammar.tokens.items())
 
     keywords = []
     grammar_labels = []
-    for label_index, new_state in arcs:
+
-        add_results(label_index)
+    def scan_stack(index):
+        dfa, state, node = stack[index]
+        states, first = dfa
+        arcs = states[state]
+
+        for label_index, new_state in arcs:
+            if label_index == 0:
+                # An accepting state, check the stack below.
+                scan_stack(index - 1)
+            else:
+                add_results(label_index)
+
+    scan_stack(-1)
 
     return keywords, grammar_labels
 

jedi/parser/__init__.py

@@ -59,6 +59,8 @@ def load_grammar(version='3.4'):
 
 
 class ErrorStatement(object):
+    type = 'error_stmt'
+
     def __init__(self, stack, arcs, next_token, position_modifier, next_start_pos):
         self.stack = stack
         self.arcs = arcs
@@ -104,15 +106,6 @@ class ErrorStatement(object):
         first = next(iterator)
         return first.get_code(include_prefix=include_prefix) + ''.join(node.get_code() for node in iterator)
 
-    def get_next_leaf(self):
-        for child in self.parent.children:
-            if child.start_pos == self.end_pos:
-                return child.first_leaf()
-
-        if child.start_pos > self.end_pos:
-            raise NotImplementedError('Node not found, must be in error statements.')
-        raise ValueError("Doesn't have a next leaf")
-
     def set_parent(self, root_node):
         """
         Used by the parser at the end of parsing. The error statements parents
@@ -313,14 +306,14 @@ class Parser(object):
             endmarker._start_pos = endmarker._start_pos[0] - 1, len(last_line)
         else:
             try:
-                newline = endmarker.get_previous()
+                newline = endmarker.get_previous_leaf()
             except IndexError:
                 return  # This means that the parser is empty.
             while True:
                 if newline.value == '':
                     # Must be a DEDENT, just continue.
                     try:
-                        newline = newline.get_previous()
+                        newline = newline.get_previous_leaf()
                     except IndexError:
                         # If there's a statement that fails to be parsed, there
                         # will be no previous leaf. So just ignore it.

jedi/parser/tree.py

@@ -221,6 +221,52 @@ class Base(object):
 
         return '\n'.join(lines)
 
+    def get_previous_leaf(self):
+        """
+        Returns the previous leaf in the parser tree.
+        Raises an IndexError if it's the first element.
+        """
+        node = self
+        while True:
+            c = node.parent.children
+            i = c.index(node)
+            if i == 0:
+                node = node.parent
+                if node.parent is None:
+                    raise IndexError('Cannot access the previous element of the first one.')
+            else:
+                node = c[i - 1]
+                break
+
+        while True:
+            try:
+                node = node.children[-1]
+            except AttributeError:  # A Leaf doesn't have children.
+                return node
+
+    def get_next_leaf(self):
+        """
+        Returns the previous leaf in the parser tree.
+        Raises an IndexError if it's the last element.
+        """
+        node = self
+        while True:
+            c = node.parent.children
+            i = c.index(node)
+            if i == len(c) - 1:
+                node = node.parent
+                if node.parent is None:
+                    raise IndexError('Cannot access the next element of the last one.')
+            else:
+                node = c[i + 1]
+                break
+
+        while True:
+            try:
+                node = node.children[0]
+            except AttributeError:  # A Leaf doesn't have children.
+                return node
+
 
 class Leaf(Base):
     __slots__ = ('position_modifier', 'value', 'parent', '_start_pos', 'prefix')
@@ -242,7 +288,7 @@ class Leaf(Base):
 
     def get_start_pos_of_prefix(self):
         try:
-            return self.get_previous().end_pos
+            return self.get_previous_leaf().end_pos
         except IndexError:
             return 1, 0  # It's the first leaf.
 
@@ -255,57 +301,9 @@ class Leaf(Base):
         self._start_pos = (self._start_pos[0] + line_offset,
                            self._start_pos[1] + column_offset)
 
-    def get_previous(self):
-        """
-        Returns the previous leaf in the parser tree.
-        Raises an IndexError if it's the first element.
-        # TODO rename to get_previous_leaf
-        """
-        node = self
-        while True:
-            c = node.parent.children
-            i = c.index(self)
-            if i == 0:
-                node = node.parent
-                if node.parent is None:
-                    raise IndexError('Cannot access the previous element of the first one.')
-            else:
-                node = c[i - 1]
-                break
-
-        while True:
-            try:
-                node = node.children[-1]
-            except AttributeError:  # A Leaf doesn't have children.
-                return node
-
     def first_leaf(self):
         return self
 
-    def get_next_leaf(self):
-        """
-        Returns the previous leaf in the parser tree.
-        Raises an IndexError if it's the last element.
-        """
-        node = self
-        while True:
-            c = node.parent.children
-            i = c.index(self)
-            if i == len(c) - 1:
-                node = node.parent
-                if node.parent is None:
-                    raise IndexError('Cannot access the next element of the last one.')
-            else:
-                node = c[i + 1]
-                break
-
-        while True:
-            try:
-                node = node.children[0]
-            except AttributeError:  # A Leaf doesn't have children.
-                return node
-
-
     def get_code(self, normalized=False, include_prefix=True):
         if normalized:
             return self.value
@@ -541,13 +539,13 @@ class BaseNode(Base):
     def get_leaf_for_position(self, position, include_prefixes=False):
         for c in self.children:
             if include_prefixes:
-                start_pos = c.get_start_pos_with_prefix()
+                start_pos = c.get_start_pos_of_prefix()
             else:
                 start_pos = c.start_pos
 
             if start_pos <= position <= c.end_pos:
                 try:
-                    return c.get_leaf_for_position(position)
+                    return c.get_leaf_for_position(position, include_prefixes)
                 except AttributeError:
                     return c
 

test/completion/on_import.py

@@ -63,8 +63,10 @@ import datetime.date
 
 #? 21 ['import']
 from import_tree.pkg import pkg
-#? ['mod1', 'mod2', 'random', 'pkg', 'rename1', 'rename2', 'recurse_class1', 'recurse_class2', 'invisible_pkg', 'flow_import']
+#? 49 ['a', '__name__', '__doc__', '__file__', '__package__']
-from import_tree.pkg import pkg, 
+from import_tree.pkg.mod1 import not_existant,    # whitespace before
+#? ['a', '__name__', '__doc__', '__file__', '__package__']
+from import_tree.pkg.mod1 import not_existant, 
 #? 22 ['mod1']
 from import_tree.pkg. import mod1
 #? 17 ['mod1', 'mod2', 'random', 'pkg', 'rename1', 'rename2', 'recurse_class1', 'recurse_class2', 'invisible_pkg', 'flow_import']