From c9fa335145204ebc5b2ed5b8c097e64f484c2aaf Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Mon, 29 Jan 2018 08:58:59 +0100
Subject: [PATCH] Fix a goto_assignments issue with a better internal API

Fixes #996.
---
 jedi/api/__init__.py               | 4 +---
 jedi/evaluate/filters.py           | 7 +++++++
 test/test_evaluate/test_imports.py | 5 +++++
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/jedi/api/__init__.py b/jedi/api/__init__.py
index 6fc1fa2f..92d0b43e 100644
--- a/jedi/api/__init__.py
+++ b/jedi/api/__init__.py
@@ -212,9 +212,7 @@ class Script(object):
 
         if follow_imports:
             def check(name):
-                if isinstance(name, ModuleName):
-                    return False
-                return name.api_type == 'module'
+                return name.is_import()
         else:
             def check(name):
                 return isinstance(name, imports.SubModuleName)
diff --git a/jedi/evaluate/filters.py b/jedi/evaluate/filters.py
index 88e4a62c..a5cce241 100644
--- a/jedi/evaluate/filters.py
+++ b/jedi/evaluate/filters.py
@@ -42,6 +42,9 @@ class AbstractNameDefinition(object):
     def execute_evaluated(self, *args, **kwargs):
         return self.infer().execute_evaluated(*args, **kwargs)
 
+    def is_import(self):
+        return False
+
     @property
     def api_type(self):
         return self.parent_context.api_type
@@ -55,6 +58,10 @@ class AbstractTreeName(AbstractNameDefinition):
     def goto(self):
         return self.parent_context.evaluator.goto(self.parent_context, self.tree_name)
 
+    def is_import(self):
+        imp = search_ancestor(self.tree_name, 'import_from', 'import_name')
+        return imp is not None
+
     @property
     def string_name(self):
         return self.tree_name.value
diff --git a/test/test_evaluate/test_imports.py b/test/test_evaluate/test_imports.py
index 74069f10..5ffea107 100644
--- a/test/test_evaluate/test_imports.py
+++ b/test/test_evaluate/test_imports.py
@@ -199,6 +199,11 @@ def test_goto_following_on_imports(Script):
     assert (g[0].line, g[0].column) != (0, 0)
 
 
+def test_goto_assignments(Script):
+    sys, = Script("import sys", 1, 10).goto_assignments(follow_imports=True)
+    assert sys.type == 'module'
+
+
 def test_os_after_from(Script):
     def check(source, result, column=None):
         completions = Script(source, column=column).completions()