Allow Jedi "goto" to perform multiple hops for "go to definition"

pylsp/plugins/definition.py

@@ -1,22 +1,54 @@
 # Copyright 2017-2020 Palantir Technologies, Inc.
 # Copyright 2021- Python Language Server Contributors.
-
+from __future__ import annotations
 import logging
+from typing import Any, Dict, List, TYPE_CHECKING
 from pylsp import hookimpl, uris, _utils
 
+if TYPE_CHECKING:
+    from jedi.api import Script
+    from jedi.api.classes import Name
+    from pylsp.config.config import Config
+    from pylsp.workspace import Document
+
 log = logging.getLogger(__name__)
 
 
+MAX_JEDI_GOTO_HOPS = 100
+
+
+def _resolve_definition(
+    maybe_defn: Name, script: Script, settings: Dict[str, Any]
+) -> Name:
+    for _ in range(MAX_JEDI_GOTO_HOPS):
+        if maybe_defn.is_definition() or maybe_defn.module_path != script.path:
+            break
+        defns = script.goto(
+            follow_imports=settings.get("follow_imports", True),
+            follow_builtin_imports=settings.get("follow_builtin_imports", True),
+            line=maybe_defn.line,
+            column=maybe_defn.column,
+        )
+        if len(defns) == 1:
+            maybe_defn = defns[0]
+        else:
+            break
+    return maybe_defn
+
+
 @hookimpl
-def pylsp_definitions(config, document, position):
+def pylsp_definitions(
+    config: Config, document: Document, position: Dict[str, int]
+) -> List[Dict[str, Any]]:
     settings = config.plugin_settings("jedi_definition")
     code_position = _utils.position_to_jedi_linecolumn(document, position)
-    definitions = document.jedi_script(use_document_path=True).goto(
+    script = document.jedi_script(use_document_path=True)
+    definitions = script.goto(
         follow_imports=settings.get("follow_imports", True),
         follow_builtin_imports=settings.get("follow_builtin_imports", True),
         **code_position,
     )
-
+    definitions = [_resolve_definition(d, script, settings) for d in definitions]
     follow_builtin_defns = settings.get("follow_builtin_definitions", True)
     return [
         {
@@ -31,7 +63,7 @@ def pylsp_definitions(config, document, position):
     ]
 
 
-def _not_internal_definition(definition):
+def _not_internal_definition(definition: Name) -> bool:
     return (
         definition.line is not None
         and definition.column is not None

test/plugins/test_definitions.py

@@ -12,7 +12,7 @@
 DOC = """def a():
     pass
 
-print a()
+print(a())
 
 
 class Directory(object):
@@ -21,6 +21,21 @@ def __init__(self):
 
     def add_member(self, id, name):
         self.members[id] = name
+
+
+subscripted_before_reference = {}
+subscripted_before_reference[0] = 0
+subscripted_before_reference
+
+
+def my_func():
+    print('called')
+
+alias = my_func
+my_list = [1, None, alias]
+inception = my_list[2]
+
+inception()
 """
 
 
@@ -40,6 +55,40 @@ def test_definitions(config, workspace):
     )
 
 
+def test_indirect_definitions(config, workspace):
+    # Over 'subscripted_before_reference'
+    cursor_pos = {"line": 16, "character": 0}
+
+    # The definition of 'subscripted_before_reference',
+    # skipping intermediate writes to the most recent definition
+    def_range = {
+        "start": {"line": 14, "character": 0},
+        "end": {"line": 14, "character": len("subscripted_before_reference")},
+    }
+
+    doc = Document(DOC_URI, workspace, DOC)
+    assert [{"uri": DOC_URI, "range": def_range}] == pylsp_definitions(
+        config, doc, cursor_pos
+    )
+
+
+def test_definition_with_multihop_inference_goto(config, workspace):
+    # Over 'inception()'
+    cursor_pos = {"line": 26, "character": 0}
+
+    # The most recent definition of 'inception',
+    # ignoring alias hops
+    def_range = {
+        "start": {"line": 24, "character": 0},
+        "end": {"line": 24, "character": len("inception")},
+    }
+
+    doc = Document(DOC_URI, workspace, DOC)
+    assert [{"uri": DOC_URI, "range": def_range}] == pylsp_definitions(
+        config, doc, cursor_pos
+    )
+
+
 def test_builtin_definition(config, workspace):
     # Over 'i' in dict
     cursor_pos = {"line": 8, "character": 24}