Further import completion improvements.

jedi/api/completion.py

@@ -58,7 +58,7 @@ class Completion:
     def __init__(self, evaluator, parser, user_context, position, call_signatures_method):
         self._evaluator = evaluator
         self._parser = parser
-        self._module = parser.module()
+        self._module = evaluator.wrap(parser.module())
         self._user_context = user_context
         self._pos = position
         self._call_signatures_method = call_signatures_method
@@ -136,7 +136,7 @@ class Completion:
         allowed_keywords, allowed_tokens = \
             helpers.get_possible_completion_types(grammar, stack)
 
-        print(allowed_keywords, allowed_tokens)
+        print(allowed_keywords, [token.tok_name[a] for a in allowed_tokens])
         completion_names = list(self._get_keyword_completion_names(allowed_keywords))
 
         if token.NAME in allowed_tokens:
@@ -145,12 +145,39 @@ class Completion:
             symbol_names = list(stack.get_node_names(grammar))
             print(symbol_names)
 
+            nodes = list(stack.get_nodes())
+            last_symbol = symbol_names[-1]
+
             if "import_stmt" in symbol_names:
-                if symbol_names[-1] == "dotted_name":
-                    completion_names += self._complete_dotted_name(stack, self._module)
-                elif symbol_names[-1] == "import_name":
-                    names = list(stack.get_nodes())[1::2]
-                    completion_names += self._get_importer_names(names)
+                level = 0
+                only_modules = True
+                level, names = self._parse_dotted_names(nodes)
+                if "import_from" in symbol_names:
+                    if 'import' in nodes:
+                        only_modules = False
+                    '''
+                    if last_symbol == "dotted_name":
+                    elif last_symbol == "import_from":
+                        # No names are given yet, but the dots for level might be
+                        # there.
+                        if 'import' in nodes:
+                            print(nodes[1])
+                            raise NotImplementedError
+                        else:
+                            raise NotImplementedError
+                    elif last_symbol == "import_name":
+                        names = nodes[1::2]
+                        completion_names += self._get_importer_names(names)
+                    '''
+                else:
+                    assert "import_name" in symbol_names
+
+                print(names, level)
+                completion_names += self._get_importer_names(
+                    names,
+                    level,
+                    only_modules
+                )
             else:
                 completion_names += self._simple_complete(completion_parts)
 
@@ -238,18 +265,20 @@ class Completion:
                 )
         return completion_names
 
-    def _complete_dotted_name(self, stack, module):
-        nodes = list(stack.get_nodes())
-
+    def _parse_dotted_names(self, nodes):
         level = 0
-        for i, node in enumerate(nodes[1:], 1):
+        names = []
+        for node in nodes[1:]:
             if node in ('.', '...'):
-                level += len(node.value)
+                if not names:
+                    level += len(node.value)
+            elif node.type == 'dotted_name':
+                names += node.children[::2]
+            elif node.type == 'name':
+                names.append(node)
             else:
-                names = nodes[i::2]
                 break
-
-        return self._get_importer_names(names, level=level)
+        return level, names
 
     def _get_importer_names(self, names, level=0, only_modules=True):
         names = [str(n) for n in names]

jedi/api/helpers.py

@@ -73,7 +73,7 @@ def get_stack_at_position(grammar, module, pos):
     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:
+    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 = ''
@@ -81,16 +81,23 @@ def get_stack_at_position(grammar, module, pos):
         for error_statement in module.error_statements:
             if error_statement.start_pos < pos <= error_statement.end_pos:
                 code = error_statement.get_code(include_prefix=False)
-                start_pos = error_statement.start_pos
+                node = error_statement
                 break
+        else:
+            raise NotImplementedError
     else:
         code = user_stmt.get_code_with_error_statements(include_prefix=False)
-        start_pos = user_stmt.start_pos
+        node = user_stmt
 
-    # Remove indentations.
-    code = code.lstrip()
-    code = get_code_until(code, start_pos, pos)
-    # Remove whitespace at the end.
+    # 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)
+    # 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()
 
     class EndMarkerReached(Exception):
@@ -101,7 +108,6 @@ def get_stack_at_position(grammar, module, pos):
             if token_[0] == token.ENDMARKER:
                 raise EndMarkerReached()
             else:
-                print(token_, token.tok_name[token_[0]])
                 yield token_
 
     print(repr(code))