From ebd080a0fd018eb326f73473303fa10515e201bd Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Wed, 3 Aug 2016 18:05:08 +0200
Subject: [PATCH] Implement goto_assignments(follow_imports=True). Fixes #382.

---
 jedi/api/__init__.py              | 21 +++++++++++++++++----
 sith.py                           |  6 +++++-
 test/test_api/test_api.py         | 13 +++++++++++++
 test/test_api/test_interpreter.py |  4 ++--
 4 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/jedi/api/__init__.py b/jedi/api/__init__.py
index b9070493..5e0d6d4e 100644
--- a/jedi/api/__init__.py
+++ b/jedi/api/__init__.py
@@ -196,7 +196,7 @@ class Script(object):
         # the API.
         return helpers.sorted_definitions(set(defs))
 
-    def goto_assignments(self):
+    def goto_assignments(self, follow_imports=False):
         """
         Return the first definition found. Imports and statements aren't
         followed. Multiple objects may be returned, because Python itself is a
@@ -205,9 +205,22 @@ class Script(object):
 
         :rtype: list of :class:`classes.Definition`
         """
-        results = self._goto()
-        d = [classes.Definition(self._evaluator, d) for d in set(results)]
-        return helpers.sorted_definitions(d)
+        def filter_follow_imports(names):
+            for name in names:
+                definition = name.get_definition()
+                if definition.type in ('import_name', 'import_from'):
+                    imp = imports.ImportWrapper(self._evaluator, name)
+                    for name in filter_follow_imports(imp.follow(is_goto=True)):
+                        yield name
+                else:
+                    yield name
+
+        names = self._goto()
+        if follow_imports:
+            names = filter_follow_imports(names)
+
+        defs = [classes.Definition(self._evaluator, d) for d in set(names)]
+        return helpers.sorted_definitions(defs)
 
     def _goto(self):
         """
diff --git a/sith.py b/sith.py
index 8b9996cc..903ef0ef 100755
--- a/sith.py
+++ b/sith.py
@@ -118,7 +118,11 @@ class TestCase(object):
         try:
             with open(self.path) as f:
                 self.script = jedi.Script(f.read(), self.line, self.column, self.path)
-            self.objects = getattr(self.script, self.operation)()
+            kwargs = {}
+            if self.operation == 'goto_assignments':
+                kwargs['follow_imports'] = random.choice([False, True])
+
+            self.objects = getattr(self.script, self.operation)(**kwargs)
             if print_result:
                 print("{path}: Line {line} column {column}".format(**self.__dict__))
                 self.show_location(self.line, self.column)
diff --git a/test/test_api/test_api.py b/test/test_api/test_api.py
index a4665f78..d388c517 100644
--- a/test/test_api/test_api.py
+++ b/test/test_api/test_api.py
@@ -164,3 +164,16 @@ def test_get_line_code():
     assert get_line_code(source, line=2) == line
     assert get_line_code(source, line=2, after=1) == line + '\nother_line'
     assert get_line_code(source, line=2, after=1, before=1) == source
+
+
+def test_goto_assignments_follow_imports():
+    code = dedent("""
+    import inspect
+    inspect.isfunction""")
+    definition, = api.Script(code, column=0).goto_assignments(follow_imports=True)
+    assert 'inspect.py' in definition.module_path
+    assert definition.start_pos == (1, 0)
+
+    definition, = api.Script(code).goto_assignments(follow_imports=True)
+    assert 'inspect.py' in definition.module_path
+    assert definition.start_pos > (1, 0)
diff --git a/test/test_api/test_interpreter.py b/test/test_api/test_interpreter.py
index 73773db1..0976241f 100644
--- a/test/test_api/test_interpreter.py
+++ b/test/test_api/test_interpreter.py
@@ -153,8 +153,8 @@ class TestInterpreterAPI(TestCase):
         def foo(bar):
             pass
 
-        lambd = lambda x: 3
+        lambd = lambda xyz: 3
 
         self.check_interpreter_complete('foo(bar', locals(), ['bar'])
         # TODO we're not yet using the Python3.5 inspect.signature, yet.
-        assert not jedi.Interpreter('lambd(x', [locals()]).completions()
+        assert not jedi.Interpreter('lambd(xyz', [locals()]).completions()