Fix a follow_imports (goto) issue.

jedi/api/__init__.py

@@ -223,23 +223,22 @@ def goto_assignments(self, follow_imports=False):
 
         :rtype: list of :class:`classes.Definition`
         """
-        def filter_follow_imports(names, follow_classes):
+        def filter_follow_imports(names, check):
             for name in names:
-                if isinstance(name, follow_classes):
+                if check(name):
                     for context in name.infer():
                         yield context.name
                 else:
                     yield name
 
         names = self._goto()
         if follow_imports:
-            # TODO really, sure? TreeNameDefinition? Should probably not follow
-            # that.
-            follow_classes = (imports.ImportName, TreeNameDefinition)
+            names = filter_follow_imports(names, lambda name: name.api_type == 'module')
         else:
-            follow_classes = (imports.SubModuleName,)
-
-        names = filter_follow_imports(names, follow_classes)
+            names = filter_follow_imports(
+                names,
+                lambda name: isinstance(name, imports.SubModuleName)
+            )
 
         defs = [classes.Definition(self._evaluator, d) for d in set(names)]
         return helpers.sorted_definitions(defs)

test/test_api/test_api.py

@@ -207,6 +207,9 @@ def test_goto_assignments_follow_imports():
     definition, = script.goto_assignments()
     assert (definition.line, definition.column) == start_pos
 
+    d, = api.Script('a = 1\na').goto_assignments(follow_imports=True)
+    assert d.name == 'a'
+
 
 def test_goto_module():
     def check(line, expected):