Repository: python-lsp/python-lsp-server
Pull request #443 (2023-09-15 16:14:05)

Allow Jedi "goto" to perform multiple hops for "go to definition"
The implementation of lsp definitions performs one hop of jedi "goto". This is inadequate for cases such as the following:
```python
d = {}
d[0]
d
```
A "definitions" request on the final reference to `d` will run a jedi goto that finds the reference to d in `d[0] = 0`, which then gets filtered out by the language server because it's not considered a definition. To fix, when we run a `goto` that finds something that's not a definition, keep running `goto`s on the thing we found until we either do find one, or exhaust the chain.

Added a unit test to cover this case, as well as type definitions to the definitions plugin.

ccordoba12: @krassowski, this looks good to me now. Are you ok with merging?
krassowski: Yes. My only suggestion would be adding another test based on the code snippet in https://jedi.readthedocs.io/en/latest/docs/api.html#type-inference-goto to show that this still works as expected.
krassowski: Also, was there ever a discussion upstream on why this is preferred behaviour in jedi?
ccordoba12: > Also, was there ever a discussion upstream on why this is preferred behaviour in jedi?

Not to my knowledge. @davidhalter, we'd appreciate your input here if you have some free time.
davidhalter: This should probably be changed upstream, yes. I feel like this implementation has almost no upsides and a lot of downsides when it comes to potential caching issues or recursive stuff (it's probably also considerably slower).
smacke: > This should probably be changed upstream, yes.

Depending on the application, I can see the value of a `goto` that visits intermediate non-definition writes and doesn't go straight to the definition. If this were changed upstream, I think it would make most sense to add a new kwarg to jedi's `goto` called `skip_to_definition` or something like that and have it default to `False` to avoid breaking anything that relies on the existing behavior. Then in pylsp we would set `skip_to_definition=True`.

I unfortunately won't have cycles for a little while to implement anything on the jedi side if the consensus is to change this upstream. If other folks are similarly lacking cycles, maybe we could add something in pylsp that does the following:

- if the signature of `goto` supports `skip_to_definition` or whatever, use that
- otherwise, use the strategy in this PR that does multiple hops

That way, whenever somebody does have cycles to submit a PR to jedi implementing this suggestion (assuming it makes sense of course), it will get automatically picked up by pylsp.

^ What do folks think about that?
ccordoba12: > If other folks are similarly lacking cycles

Yeah, I think we don't have time for that.

> maybe we could add something in pylsp that does the following:
>
> * if the signature of goto supports skip_to_definition or whatever, use that
> * otherwise, use the strategy in this PR that does multiple hops

Until this is fixed in Jedi, I think there's no need to implement the first part and only go with the second one (which is what you've already done).

@krassowski, what do you think?
krassowski: Yes, we are version pinning jedi in python-lsp-server anyways so it makes sense to keep it as-is. Let's just add the test case I mentioned earlier, open an issue to track the future planned improvement (whether on jedi side or in here)  and then I think we can merge this patch in.
smacke: > Let's just add the test case I mentioned earlier

Done! (My assumption is that we do **not** want the multihop behavior in this case, since the 1-hop goto is still an actual definition)

> open an issue to track the future planned improvement (whether on jedi side or in here)

I will do it here: https://github.com/python-lsp/python-lsp-server/issues/446

Patch:
diff --git a/pylsp/plugins/definition.py b/pylsp/plugins/definition.py
index a5ccbd7..ffc1b00 100644
--- a/pylsp/plugins/definition.py
+++ b/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 patch:
diff --git a/test/plugins/test_definitions.py b/test/plugins/test_definitions.py
index 34acc6a..f0e9ffe 100644
--- a/test/plugins/test_definitions.py
+++ b/test/plugins/test_definitions.py
@@ -12,7 +12,7 @@ DOC_URI = uris.from_fs_path(__file__)
 DOC = """def a():
     pass
 
-print a()
+print(a())
 
 
 class Directory(object):
@@ -21,6 +21,21 @@ class Directory(object):
 
     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}