again testing improvements. follow_definitions imports are now pretty much covered

jedi/api/classes.py

@@ -400,15 +400,18 @@ class Completion(BaseDefinition):
         if self._definition.isinstance(pr.Statement):
             defs = self._evaluator.eval_statement(self._definition)
         elif self._definition.isinstance(pr.Import):
-            print('ha', self._definition, repr(self._name))
-            i = imports.ImportPath(self._evaluator, self._definition, True)
-            print('h', i.import_path + [unicode(self._name)])
-            defs = imports.Importer(i.import_path + [unicode(self._name)],
-                                    i._importer.module).follow_file_system()[0]
-            defs = [defs]
-            print(defs)
-            #defs = imports.strip_imports(self._evaluator, [self._definition])
+            if self._definition.alias is None:
+                print('ha', self._definition, repr(self._name))
+                i = imports.ImportPath(self._evaluator, self._definition, True)
+                print('h', i.import_path + [unicode(self._name)])
+                defs = imports.Importer(i.import_path + [unicode(self._name)],
+                                        i._importer.module).follow_file_system()[0]
+                defs = [defs]
+                print(defs)
+            else:
+                defs = imports.strip_imports(self._evaluator, [self._definition])
         else:
+            print('else', self._definition)
             return [self]
 
         defs = [BaseDefinition(self._evaluator, d, d.start_pos) for d in defs]

jedi/evaluate/imports.py

@@ -266,6 +266,31 @@ class Importer(object):
 
         return in_path + sys_path.sys_path_with_modifications(self.module)
 
+    def follow(self, evaluator):
+        try:
+            scope, rest = self._importer.follow_file_system()
+        except ModuleNotFound:
+            debug.warning('Module not found: %s', self.import_stmt)
+            return []
+
+        scopes = [scope]
+        scopes += remove_star_imports(self._evaluator, scope)
+
+        # follow the rest of the import (not FS -> classes, functions)
+        if len(rest) > 1 or rest and self.is_like_search:
+            scopes = []
+            if ['os', 'path'] == self.import_path[:2] \
+                    and not self._is_relative_import():
+                # This is a huge exception, we follow a nested import
+                # ``os.path``, because it's a very important one in Python
+                # that is being achieved by messing with ``sys.modules`` in
+                # ``os``.
+                scopes = self._evaluator.follow_path(iter(rest), [scope], scope)
+        elif rest:
+            scopes = itertools.chain.from_iterable(
+                self._evaluator.follow_path(iter(rest), [s], s)
+                for s in scopes)
+
     def follow_file_system(self):
         if self.file_path:
             sys_path_mod = list(self.sys_path_with_modifications())

test/test_evaluate/test_imports.py

@@ -1,3 +1,5 @@
+from itertools import chain
+
 import pytest
 
 import jedi
@@ -33,24 +35,50 @@ def test_import_empty():
     assert definition
 
 
+def check_follow_definition_types(source):
+    # nested import
+    completions = jedi.Script(source, path='some_path.py').completions()
+    defs = chain.from_iterable(c.follow_definition() for c in completions)
+    return [d.type for d in defs]
+
+
 def test_follow_import_incomplete():
     """
     Completion on incomplete imports should always take the full completion
     to do any evaluation.
     """
-    datetime = jedi.Script("import datetim").completions()[0]
-    definition = datetime.follow_definition()[0]
-    assert definition
+    datetime = check_follow_definition_types("import itertool")
+    assert datetime == ['module']
 
     # empty `from * import` parts
-    datetime = jedi.Script("from datetime import ").completions()[0]
-    definitions = datetime.follow_definition()
-    assert [d.type for d in definitions if d.name == 'date'] == ['class']
+    itert = jedi.Script("from itertools import ").completions()
+    definitions = [d for d in itert if d.name == 'chain']
+    assert len(definitions) == 1
+    assert [d.type for d in definitions[0].follow_definition()] == ['class']
 
     # incomplete `from * import` part
-    datetime = jedi.Script("from datetime import datetim").completions()[0]
-    definition = datetime.follow_definition()
-    assert [d.type for d in definitions] == ['class']
+    datetime = check_follow_definition_types("from datetime import datetim")
+    assert set(datetime) == set(['class'])  # py33: builtin and pure py version
+
+    # os.path check
+    ospath = check_follow_definition_types("from os.path import abspat")
+    assert ospath == ['function']
+
+    # alias
+    alias = check_follow_definition_types("import io as abcd; abcd")
+    assert alias == ['module']
+
+
+@cwd_at('test/completion/import_tree')
+def test_follow_definition_nested_import():
+    types = check_follow_definition_types("import pkg.mod1; pkg")
+    assert types == ['module']
+
+    types = check_follow_definition_types("import pkg.mod1; pkg.mod1")
+    assert types == ['module']
+
+    types = check_follow_definition_types("import pkg.mod1; pkg.mod1.a")
+    assert types == ['class']
 
 
 def test_follow_definition_land_on_import():