Semantic Token Support

pylsp/config/schema.json

@@ -232,6 +232,11 @@
       "default": true,
       "description": "Enable or disable the plugin."
     },
+    "pylsp.plugins.semantic_tokens.enabled": {
+      "type": "boolean",
+      "default": true,
+      "description": "Enable or disable the plugin."
+    },
     "pylsp.plugins.jedi_references.enabled": {
       "type": "boolean",
       "default": true,
@@ -500,4 +505,4 @@
       "description": "The name of the folder in which rope stores project configurations and data.  Pass `null` for not using such a folder at all."
     }
   }
-}
+}

pylsp/hookspecs.py

@@ -93,6 +93,11 @@ def pylsp_hover(config, workspace, document, position):
     pass
 
 
+@hookspec(firstresult=True)
+def pylsp_semantic_tokens(config, workspace, document):
+    pass
+
+
 @hookspec
 def pylsp_initialize(config, workspace):
     pass

pylsp/lsp.py

@@ -6,6 +6,9 @@
 https://github.com/Microsoft/language-server-protocol/blob/master/protocol.md
 """
 
+from enum import Enum
+from typing import NamedTuple
+
 
 class CompletionItemKind:
     Text = 1
@@ -86,6 +89,52 @@ class SymbolKind:
     Array = 18
 
 
+class SemanticToken(NamedTuple):
+    value: int
+    name: str
+
+
+class SemanticTokenType(Enum):
+    Namespace = SemanticToken(0, "namespace")
+    # represents a generic type. acts as a fallback for types which
+    # can't be mapped to a specific type like class or enum.
+    Type = SemanticToken(1, "type")
+    Class = SemanticToken(2, "class")
+    Enum = SemanticToken(3, "enum")
+    Interface = SemanticToken(4, "interface")
+    Struct = SemanticToken(5, "struct")
+    TypeParameter = SemanticToken(6, "typeParameter")
+    Parameter = SemanticToken(7, "parameter")
+    Variable = SemanticToken(8, "variable")
+    Property = SemanticToken(9, "property")
+    EnumMember = SemanticToken(10, "enumMember")
+    Event = SemanticToken(11, "event")
+    Function = SemanticToken(12, "function")
+    Method = SemanticToken(13, "method")
+    Macro = SemanticToken(14, "macro")
+    Keyword = SemanticToken(15, "keyword")
+    Modifier = SemanticToken(16, "modifier")
+    Comment = SemanticToken(17, "comment")
+    String = SemanticToken(18, "string")
+    Number = SemanticToken(19, "number")
+    Regexp = SemanticToken(20, "regexp")
+    Operator = SemanticToken(21, "operator")
+    Decorator = SemanticToken(22, "decorator")  # @since 3.17.0
+
+
+class SemanticTokenModifier(Enum):
+    Declaration = SemanticToken(0, "declaration")
+    Definition = SemanticToken(1, "definition")
+    Readonly = SemanticToken(2, "readonly")
+    Static = SemanticToken(3, "static")
+    Deprecated = SemanticToken(4, "deprecated")
+    Abstract = SemanticToken(5, "abstract")
+    Async = SemanticToken(6, "async")
+    Modification = SemanticToken(7, "modification")
+    Documentation = SemanticToken(8, "documentation")
+    DefaultLibrary = SemanticToken(9, "defaultLibrary")
+
+
 class TextDocumentSyncKind:
     NONE = 0
     FULL = 1

pylsp/plugins/semantic_tokens.py

@@ -0,0 +1,114 @@
+# Copyright 2017-2020 Palantir Technologies, Inc.
+# Copyright 2021- Python Language Server Contributors.
+
+import logging
+
+from jedi.api.classes import Name
+
+from pylsp import hookimpl
+from pylsp.config.config import Config
+from pylsp.lsp import SemanticTokenType
+from pylsp.workspace import Document
+
+log = logging.getLogger(__name__)
+
+# Valid values for type are ``module``, ``class``, ``instance``, ``function``,
+# ``param``, ``path``, ``keyword``, ``property`` and ``statement``.
+TYPE_MAP = {
+    "module": SemanticTokenType.Namespace.value.value,
+    "class": SemanticTokenType.Class.value.value,
+    # "instance": SemanticTokenType.Type.value.value,
+    "function": SemanticTokenType.Function.value.value,
+    "param": SemanticTokenType.Parameter.value.value,
+    # "path": SemanticTokenType.Type.value.value,
+    "keyword": SemanticTokenType.Keyword.value.value,
+    "property": SemanticTokenType.Property.value.value,
+    # "statement": SemanticTokenType.Variable.value.value,
+}
+
+
+def _raw_semantic_token(n: Name) -> list[int] | None:
+    """Find an appropriate semantic token for the name.
+
+    This works by looking up the definition (using jedi ``goto``) of the name and
+    matching the definition's type to one of the availabile semantic tokens. Further
+    improvements are possible by inspecting context, e.g. semantic token modifiers such
+    as ``abstract`` or ``async`` or even different tokens, e.g. ``property`` or
+    ``method``. Dunder methods may warrant special treatment/modifiers as well.
+
+    The return is a "raw" semantic token rather than a "diff." This is in the form of a
+    length 5 array of integers where the elements are the line number, starting
+    character, length, token index, and modifiers (as an integer whose binary
+    representation has bits set at the indices of all applicable modifiers).
+    """
+    definitions = n.goto(
+        follow_imports=True,
+        follow_builtin_imports=True,
+        only_stubs=False,
+        prefer_stubs=False,
+    )
+    if not definitions:
+        log.debug(
+            "no definitions found for name %s (%s:%s)", n.description, n.line, n.column
+        )
+        return None
+    if len(definitions) > 1:
+        log.debug(
+            "multiple definitions found for name %s (%s:%s)",
+            n.description,
+            n.line,
+            n.column,
+        )
+    definition, *_ = definitions
+    if (definition_type := TYPE_MAP.get(definition.type, None)) is None:
+        log.debug(
+            "no matching semantic token for name %s (%s:%s)",
+            n.description,
+            n.line,
+            n.column,
+        )
+        return None
+    return [n.line - 1, n.column, len(n.name), definition_type, 0]
+
+
+def _diff_position(
+    token_line: int, token_start_char: int, current_line: int, current_start_char: int
+) -> tuple[int, int, int, int]:
+    """Compute the diff position for a semantic token.
+
+    This returns the delta line and column as well as what should be considered the
+    "new" current line and column.
+    """
+    delta_start_char = (
+        token_start_char - current_start_char
+        if token_line == current_line
+        else token_start_char
+    )
+    delta_line = token_line - current_line
+    return (delta_line, delta_start_char, token_line, token_start_char)
+
+
+@hookimpl
+def pylsp_semantic_tokens(config: Config, document: Document):
+    # Currently unused, but leaving it here for easy adding of settings.
+    symbols_settings = config.plugin_settings("semantic_tokens")
+
+    names = document.jedi_names(all_scopes=True, definitions=True, references=True)
+    data = []
+    line, start_char = 0, 0
+    for n in names:
+        token = _raw_semantic_token(n)
+        log.debug(
+            "raw token for name %s (%s:%s): %s", n.description, n.line, n.column, token
+        )
+        if token is None:
+            continue
+        token[0], token[1], line, start_char = _diff_position(
+            token[0], token[1], line, start_char
+        )
+        log.debug(
+            "diff token for name %s (%s:%s): %s", n.description, n.line, n.column, token
+        )
+        data.extend(token)
+
+    return {"data": data}

pylsp/python_lsp.py

@@ -276,6 +276,24 @@ def capabilities(self):
                 "commands": flatten(self._hook("pylsp_commands"))
             },
             "hoverProvider": True,
+            "semanticTokensProvider": {
+                "legend": {
+                    "tokenTypes": [
+                        semantic_token_type.value.name
+                        for semantic_token_type in sorted(
+                            lsp.SemanticTokenType, key=lambda x: x.value
+                        )
+                    ],
+                    "tokenModifiers": [
+                        semantic_token_modifier.value.name
+                        for semantic_token_modifier in sorted(
+                            lsp.SemanticTokenModifier, key=lambda x: x.value
+                        )
+                    ],
+                },
+                "range": False,
+                "full": True,
+            },
             "referencesProvider": True,
             "renameProvider": True,
             "foldingRangeProvider": True,
@@ -432,6 +450,9 @@ def highlight(self, doc_uri, position):
     def hover(self, doc_uri, position):
         return self._hook("pylsp_hover", doc_uri, position=position) or {"contents": ""}
 
+    def semantic_tokens(self, doc_uri):
+        return self._hook("pylsp_semantic_tokens", doc_uri) or {"data": []}
+
     @_utils.debounce(LINT_DEBOUNCE_S, keyed_by="doc_uri")
     def lint(self, doc_uri, is_saved):
         # Since we're debounced, the document may no longer be open
@@ -760,6 +781,9 @@ def m_text_document__document_highlight(
     def m_text_document__hover(self, textDocument=None, position=None, **_kwargs):
         return self.hover(textDocument["uri"], position)
 
+    def m_text_document__semantic_tokens__full(self, textDocument=None, **_kwargs):
+        return self.semantic_tokens(textDocument["uri"])
+
     def m_text_document__document_symbol(self, textDocument=None, **_kwargs):
         return self.document_symbols(textDocument["uri"])
 

pyproject.toml

@@ -67,6 +67,7 @@ flake8 = "pylsp.plugins.flake8_lint"
 jedi_completion = "pylsp.plugins.jedi_completion"
 jedi_definition = "pylsp.plugins.definition"
 jedi_hover = "pylsp.plugins.hover"
+semantic_tokens = "pylsp.plugins.semantic_tokens"
 jedi_highlight = "pylsp.plugins.highlight"
 jedi_references = "pylsp.plugins.references"
 jedi_rename = "pylsp.plugins.jedi_rename"