From dbdd556a2be8c1735950bf3015700d8468635a7e Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Wed, 22 Jan 2020 18:29:02 +0100
Subject: [PATCH] Add follow_imports to Definition.goto, fixes #1474

---
 CHANGELOG.rst                 |  2 ++
 jedi/api/__init__.py          | 20 +-------------------
 jedi/api/classes.py           |  9 +++++++--
 jedi/api/helpers.py           | 22 ++++++++++++++++++++++
 test/test_api/test_classes.py | 14 ++++++++++++++
 5 files changed, 46 insertions(+), 21 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index e1f89a12..83eb3aa3 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -26,6 +26,8 @@ Changelog
     - ``usages`` deprecated, use ``get_references`` instead
     - ``jedi.names`` deprecated, use ``jedi.Script(...).get_names()``
 - ``BaseDefinition.goto_assignments`` renamed to ``BaseDefinition.goto``
+- Add follow_imports to ``Definition.goto``. Now its signature matches
+  ``Script.goto``.
 - Python 2 support deprecated. For this release it is best effort. Python 2 has
   reached the end of its life and now it's just about a smooth transition.
 
diff --git a/jedi/api/__init__.py b/jedi/api/__init__.py
index e290ed36..44751109 100644
--- a/jedi/api/__init__.py
+++ b/jedi/api/__init__.py
@@ -286,24 +286,6 @@ class Script(object):
 
     def _goto(self, line, column, follow_imports=False, follow_builtin_imports=False,
               only_stubs=False, prefer_stubs=False):
-        def filter_follow_imports(names):
-            for name in names:
-                if name.is_import():
-                    new_names = list(filter_follow_imports(name.goto()))
-                    found_builtin = False
-                    if follow_builtin_imports:
-                        for new_name in new_names:
-                            if new_name.start_pos is None:
-                                found_builtin = True
-
-                    if found_builtin:
-                        yield name
-                    else:
-                        for new_name in new_names:
-                            yield new_name
-                else:
-                    yield name
-
         tree_name = self._module_node.get_name_of_position((line, column))
         if tree_name is None:
             # Without a name we really just want to jump to the result e.g.
@@ -328,7 +310,7 @@ class Script(object):
             names = list(name.goto())
 
         if follow_imports:
-            names = filter_follow_imports(names)
+            names = helpers.filter_follow_imports(names)
         names = convert_names(
             names,
             only_stubs=only_stubs,
diff --git a/jedi/api/classes.py b/jedi/api/classes.py
index fce8bc8e..b44e5eb8 100644
--- a/jedi/api/classes.py
+++ b/jedi/api/classes.py
@@ -20,6 +20,7 @@ from jedi.inference.gradual.conversion import convert_names, convert_values
 from jedi.inference.base_value import ValueSet
 from jedi.api.keywords import KeywordName
 from jedi.api import completion_cache
+from jedi.api.helpers import filter_follow_imports
 
 
 def _sort_names_by_start_pos(names):
@@ -368,13 +369,17 @@ class BaseDefinition(object):
         )
         return self.goto(**kwargs)
 
-    def _goto(self, only_stubs=False, prefer_stubs=False):
+    def _goto(self, follow_imports=False, follow_builtin_imports=False,
+              only_stubs=False, prefer_stubs=False):
 
         if not self._name.is_value_name:
             return []
 
+        names = self._name.goto()
+        if follow_imports:
+            names = filter_follow_imports(names, follow_builtin_imports)
         names = convert_names(
-            self._name.goto(),
+            names,
             only_stubs=only_stubs,
             prefer_stubs=prefer_stubs,
         )
diff --git a/jedi/api/helpers.py b/jedi/api/helpers.py
index 598bd3ca..fbfbdb31 100644
--- a/jedi/api/helpers.py
+++ b/jedi/api/helpers.py
@@ -170,6 +170,28 @@ def infer(inference_state, context, leaf):
     return definitions
 
 
+def filter_follow_imports(names, follow_builtin_imports=False):
+    for name in names:
+        if name.is_import():
+            new_names = list(filter_follow_imports(
+                name.goto(),
+                follow_builtin_imports=follow_builtin_imports,
+            ))
+            found_builtin = False
+            if follow_builtin_imports:
+                for new_name in new_names:
+                    if new_name.start_pos is None:
+                        found_builtin = True
+
+            if found_builtin:
+                yield name
+            else:
+                for new_name in new_names:
+                    yield new_name
+        else:
+            yield name
+
+
 class CallDetails(object):
     def __init__(self, bracket_leaf, children, position):
         ['bracket_leaf', 'call_index', 'keyword_name_str']
diff --git a/test/test_api/test_classes.py b/test/test_api/test_classes.py
index e7178790..4a17c572 100644
--- a/test/test_api/test_classes.py
+++ b/test/test_api/test_classes.py
@@ -544,3 +544,17 @@ def test_inheritance_module_path(Script, goto, code, name, file_name):
     assert func.type == 'function'
     assert func.name == name
     assert func.module_path == os.path.join(base_path, file_name)
+
+
+def test_definition_goto_follow_imports(Script):
+    dumps = Script('from json import dumps\ndumps').get_names(references=True)[-1]
+    assert dumps.description == 'dumps'
+    no_follow, = dumps.goto()
+    assert no_follow.description == 'def dumps'
+    assert no_follow.line == 1
+    assert no_follow.column == 17
+    assert no_follow.module_name == '__main__'
+    follow, = dumps.goto(follow_imports=True)
+    assert follow.description == 'def dumps'
+    assert follow.line != 1
+    assert follow.module_name == 'json'