Add textDocument/typeDefinition plugin

Author: Hoblovski

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"