Actually do goto when follow_imports is used

Fixes #945.

Author: davidhalter

jedi/api/__init__.py

@@ -226,19 +226,22 @@ def goto_assignments(self, follow_imports=False):
         def filter_follow_imports(names, check):
             for name in names:
                 if check(name):
-                    for context in name.infer():
-                        yield context.name
+                    for result in filter_follow_imports(name.goto(), check):
+                        yield result
                 else:
                     yield name
 
         names = self._goto()
         if follow_imports:
-            names = filter_follow_imports(names, lambda name: name.api_type == 'module')
+            def check(name):
+                if isinstance(name, er.ModuleName):
+                    return False
+                return name.api_type == 'module'
         else:
-            names = filter_follow_imports(
-                names,
-                lambda name: isinstance(name, imports.SubModuleName)
-            )
+            def check(name):
+                return isinstance(name, imports.SubModuleName)
+
+        names = filter_follow_imports(names, check)
 
         defs = [classes.Definition(self._evaluator, d) for d in set(names)]
         return helpers.sorted_definitions(defs)

jedi/evaluate/filters.py

@@ -20,6 +20,12 @@ class AbstractNameDefinition(object):
     def infer(self):
         raise NotImplementedError
 
+    @abstractmethod
+    def goto(self):
+        # Typically names are already definitions and therefore a goto on that
+        # name will always result on itself.
+        return set([self])
+
     def get_root_context(self):
         return self.parent_context.get_root_context()
 
@@ -44,6 +50,9 @@ def __init__(self, parent_context, tree_name):
         self.parent_context = parent_context
         self.tree_name = tree_name
 
+    def goto(self):
+        return self.parent_context.evaluator.goto(self.parent_context, self.tree_name)
+
     @property
     def string_name(self):
         return self.tree_name.value

jedi/evaluate/imports.py

@@ -141,6 +141,7 @@ def get_init_path(directory_path):
 
 class ImportName(AbstractNameDefinition):
     start_pos = (1, 0)
+    _level = 0
 
     def __init__(self, parent_context, string_name):
         self.parent_context = parent_context
@@ -151,8 +152,12 @@ def infer(self):
             self.parent_context.evaluator,
             [self.string_name],
             self.parent_context,
+            level=self._level,
         ).follow()
 
+    def goto(self):
+        return [m.name for m in self.infer()]
+
     def get_root_context(self):
         # Not sure if this is correct.
         return self.parent_context.get_root_context()
@@ -163,13 +168,7 @@ def api_type(self):
 
 
 class SubModuleName(ImportName):
-    def infer(self):
-        return Importer(
-            self.parent_context.evaluator,
-            [self.string_name],
-            self.parent_context,
-            level=1
-        ).follow()
+    _level = 1
 
 
 class Importer(object):