Add support for type definition

CONFIGURATION.md

@@ -33,6 +33,7 @@ This server can be configured using the `workspace/didChangeConfiguration` metho
 | `pylsp.plugins.jedi_completion.resolve_at_most` | `integer` | How many labels and snippets (at most) should be resolved? | `25` |
 | `pylsp.plugins.jedi_completion.cache_for` | `array` of `string` items | Modules for which labels and snippets should be cached. | `["pandas", "numpy", "tensorflow", "matplotlib"]` |
 | `pylsp.plugins.jedi_definition.enabled` | `boolean` | Enable or disable the plugin. | `true` |
+| `pylsp.plugins.type_definition.enabled` | `boolean` | Enable or disable the plugin. | `true` |
 | `pylsp.plugins.jedi_definition.follow_imports` | `boolean` | The goto call will follow imports. | `true` |
 | `pylsp.plugins.jedi_definition.follow_builtin_imports` | `boolean` | If follow_imports is True will decide if it follow builtin imports. | `true` |
 | `pylsp.plugins.jedi_definition.follow_builtin_definitions` | `boolean` | Follow builtin and extension definitions to stubs. | `true` |

pylsp/config/schema.json

@@ -225,6 +225,11 @@
       "default": true,
       "description": "Enable or disable the plugin."
     },
+    "pylsp.plugins.type_definition.enabled": {
+      "type": "boolean",
+      "default": true,
+      "description": "Enable or disable the plugin."
+    },
     "pylsp.plugins.jedi_definition.follow_imports": {
       "type": "boolean",
       "default": true,

pylsp/hookspecs.py

@@ -38,6 +38,11 @@ def pylsp_definitions(config, workspace, document, position) -> None:
     pass
 
 
+@hookspec(firstresult=True)
+def pylsp_type_definition(config, document, position):
+    pass
+
+
 @hookspec
 def pylsp_dispatchers(config, workspace) -> None:
     pass

pylsp/plugins/type_definition.py

@@ -0,0 +1,40 @@
+# Copyright 2017-2020 Palantir Technologies, Inc.
+# Copyright 2021- Python Language Server Contributors.
+
+import logging
+
+from pylsp import _utils, hookimpl
+
+
+log = logging.getLogger(__name__)
+
+
+def lsp_location(name):
+    module_path = name.module_path
+    if module_path is None or name.line is None or name.column is None:
+        return None
+    uri = module_path.as_uri()
+    return {
+        "uri": str(uri),
+        "range": {
+            "start": {"line": name.line - 1, "character": name.column},
+            "end": {"line": name.line - 1, "character": name.column + len(name.name)},
+        },
+    }
+
+
+@hookimpl
+def pylsp_type_definition(config, document, position):
+    try:
+        kwargs = _utils.position_to_jedi_linecolumn(document, position)
+        script = document.jedi_script()
+        names = script.infer(**kwargs)
+        definitions = [
+            definition
+            for definition in [lsp_location(name) for name in names]
+            if definition is not None
+        ]
+        return definitions
+    except Exception as e:
+        log.debug("Failed to run type_definition: %s", e)
+        return []

pylsp/python_lsp.py

@@ -276,6 +276,7 @@ def capabilities(self):
             "documentRangeFormattingProvider": True,
             "documentSymbolProvider": True,
             "definitionProvider": True,
+            "typeDefinitionProvider": True,
             "executeCommandProvider": {
                 "commands": flatten(self._hook("pylsp_commands"))
             },
@@ -412,6 +413,9 @@ def completion_item_resolve(self, completion_item):
     def definitions(self, doc_uri, position):
         return flatten(self._hook("pylsp_definitions", doc_uri, position=position))
 
+    def type_definition(self, doc_uri, position):
+        return self._hook("pylsp_type_definition", doc_uri, position=position)
+
     def document_symbols(self, doc_uri):
         return flatten(self._hook("pylsp_document_symbols", doc_uri))
 
@@ -762,6 +766,9 @@ def m_text_document__definition(self, textDocument=None, position=None, **_kwarg
             return self._cell_document__definition(document, position, **_kwargs)
         return self.definitions(textDocument["uri"], position)
 
+    def m_text_document__type_definition(self, textDocument=None, position=None, **_kwargs):
+        return self.type_definition(textDocument["uri"], position)
+
     def m_text_document__document_highlight(
         self, textDocument=None, position=None, **_kwargs
     ):

pyproject.toml

@@ -66,6 +66,7 @@ folding = "pylsp.plugins.folding"
 flake8 = "pylsp.plugins.flake8_lint"
 jedi_completion = "pylsp.plugins.jedi_completion"
 jedi_definition = "pylsp.plugins.definition"
+type_definition = "pylsp.plugins.type_definition"
 jedi_hover = "pylsp.plugins.hover"
 jedi_highlight = "pylsp.plugins.highlight"
 jedi_references = "pylsp.plugins.references"

test/plugins/test_type_definition.py

@@ -0,0 +1,103 @@
+# Copyright 2017-2020 Palantir Technologies, Inc.
+# Copyright 2021- Python Language Server Contributors.
+
+import os
+
+from pylsp import uris
+from pylsp.plugins.type_definition import pylsp_type_definition
+from pylsp.workspace import Document
+
+DOC_URI = uris.from_fs_path(__file__)
+DOC = """\
+from dataclasses import dataclass
+
+@dataclass
+class IntPair:
+    a: int
+    b: int
+
+def main() -> None:
+    l0 = list(1, 2)
+
+    my_pair = IntPair(a=10, b=20)
+    print(f"Original pair: {my_pair}")
+"""
+
+
+def test_type_definitions(config, workspace) -> None:
+    # Over 'IntPair' in 'main'
+    cursor_pos = {"line": 10, "character": 14}
+
+    # The definition of 'IntPair'
+    def_range = {
+        "start": {"line": 3, "character": 6},
+        "end": {"line": 3, "character": 13},
+    }
+
+    doc = Document(DOC_URI, workspace, DOC)
+    assert [{"uri": DOC_URI, "range": def_range}] == pylsp_type_definition(
+        config, doc, cursor_pos
+    )
+
+
+def test_builtin_definition(config, workspace) -> None:
+    # Over 'list' in main
+    cursor_pos = {"line": 8, "character": 9}
+
+    doc = Document(DOC_URI, workspace, DOC)
+    orig_settings = config.settings()
+
+    defns = pylsp_type_definition(config, doc, cursor_pos)
+    assert len(defns) == 1
+    assert defns[0]["uri"].endswith("builtins.pyi")
+
+
+def test_mutli_file_type_definitions(config, workspace, tmpdir) -> None:
+    # Create a dummy module out of the workspace's root_path and try to get
+    # a definition on it in another file placed next to it.
+    module_content = """\
+from dataclasses import dataclass
+
+@dataclass
+class IntPair:
+    a: int
+    b: int
+"""
+    p1 = tmpdir.join("intpair.py")
+    p1.write(module_content)
+    # The uri for intpair.py
+    module_path = str(p1)
+    module_uri = uris.from_fs_path(module_path)
+
+    # Content of doc to test type definition
+    doc_content = """\
+from intpair import IntPair
+
+def main() -> None:
+    l0 = list(1, 2)
+
+    my_pair = IntPair(a=10, b=20)
+    print(f"Original pair: {my_pair}")
+"""
+    p2 = tmpdir.join("main.py")
+    p2.write(doc_content)
+    doc_path = str(p2)
+    doc_uri = uris.from_fs_path(doc_path)
+
+    doc = Document(doc_uri, workspace, doc_content)
+
+    # The range where IntPair is defined in intpair.py
+    def_range = {
+        "start": {"line": 3, "character": 6},
+        "end": {"line": 3, "character": 13},
+    }
+
+    # The position where IntPair is called in main.py
+    cursor_pos = {"line": 5, "character": 14}
+
+    print("!URI", module_uri)
+    print("!URI", doc_uri)
+
+    assert [{"uri": module_uri, "range": def_range}] == pylsp_type_definition(
+        config, doc, cursor_pos
+    )