implement document highlight for variables and keywords

Author: d-biehl

CHANGELOG.md

@@ -7,6 +7,7 @@ All notable changes to the "robotcode" extension will be documented in this file
 ### added
 
 - implement find references for libraries, resources, variables import
+- implement document highlight for variables and keywords
 
 ##  0.4.2
 

robotcode/language_server/common/lsp_types.py

@@ -931,6 +931,16 @@ class ReferenceRegistrationOptions(TextDocumentRegistrationOptions, ReferenceOpt
     pass
 
 
+@dataclass(repr=False)
+class DocumentHighlightOptions(WorkDoneProgressOptions):
+    pass
+
+
+@dataclass(repr=False)
+class DocumentHighlightRegistrationOptions(TextDocumentRegistrationOptions, DocumentHighlightOptions):
+    pass
+
+
 @dataclass(repr=False)
 class ServerCapabilities(Model):
     text_document_sync: Union[TextDocumentSyncOptions, TextDocumentSyncKind, None] = None
@@ -941,7 +951,7 @@ class ServerCapabilities(Model):
     definition_provider: Union[bool, DefinitionOptions, None] = None
     implementation_provider: Union[bool, ImplementationOptions, ImplementationRegistrationOptions, None] = None
     references_provider: Union[bool, ReferenceOptions, None] = None
-    # TODO document_highlight_provider: Union[bool, DocumentHighlightOptions, None] = None
+    document_highlight_provider: Union[bool, DocumentHighlightOptions, None] = None
     document_symbol_provider: Union[bool, DocumentSymbolOptions, None] = None
     # TODO code_action_provider: Union[bool, CodeActionOptions] = None
     code_lens_provider: Optional[CodeLensOptions] = None
@@ -1837,3 +1847,27 @@ class _ReferenceParams(Model):
 @dataclass(repr=False)
 class ReferenceParams(WorkDoneProgressParams, PartialResultParams, TextDocumentPositionParams, _ReferenceParams):
     pass
+
+
+@dataclass(repr=False)
+class _DocumentHighlightParams(Model):
+    pass
+
+
+@dataclass(repr=False)
+class DocumentHighlightParams(
+    WorkDoneProgressParams, PartialResultParams, TextDocumentPositionParams, _DocumentHighlightParams
+):
+    pass
+
+
+class DocumentHighlightKind(Enum):
+    TEXT = 1
+    READ = 2
+    WRITE = 3
+
+
+@dataclass(repr=False)
+class DocumentHighlight(Model):
+    range: Range
+    kind: Optional[DocumentHighlightKind] = None

robotcode/language_server/common/parts/document_highlight.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from asyncio import CancelledError
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from ....jsonrpc2.protocol import rpc_method
+from ....utils.async_tools import async_tasking_event
+from ....utils.logging import LoggingDescriptor
+from ..has_extend_capabilities import HasExtendCapabilities
+from ..language import language_id_filter
+from ..lsp_types import (
+    DocumentHighlight,
+    DocumentHighlightOptions,
+    DocumentHighlightParams,
+    Position,
+    ServerCapabilities,
+    TextDocumentIdentifier,
+)
+from ..text_document import TextDocument
+
+if TYPE_CHECKING:
+    from ..protocol import LanguageServerProtocol
+
+from .protocol_part import LanguageServerProtocolPart
+
+
+class DocumentHighlightProtocolPart(LanguageServerProtocolPart, HasExtendCapabilities):
+
+    _logger = LoggingDescriptor()
+
+    def __init__(self, parent: LanguageServerProtocol) -> None:
+        super().__init__(parent)
+
+    def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
+        if len(self.collect):
+            capabilities.document_highlight_provider = DocumentHighlightOptions(work_done_progress=True)
+
+    @async_tasking_event
+    async def collect(
+        sender, document: TextDocument, position: Position  # NOSONAR
+    ) -> Optional[List[DocumentHighlight]]:
+        ...
+
+    @rpc_method(name="textDocument/documentHighlight", param_type=DocumentHighlightParams)
+    async def _text_document_document_highlight(
+        self,
+        text_document: TextDocumentIdentifier,
+        position: Position,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Optional[List[DocumentHighlight]]:
+
+        highlights: List[DocumentHighlight] = []
+
+        document = await self.parent.documents.get(text_document.uri)
+        if document is None:
+            return None
+
+        for result in await self.collect(self, document, position, callback_filter=language_id_filter(document)):
+            if isinstance(result, BaseException):
+                if not isinstance(result, CancelledError):
+                    self._logger.exception(result, exc_info=result)
+            else:
+                if result is not None:
+                    highlights.extend(result)
+
+        if len(highlights) == 0:
+            return None
+
+        return highlights

robotcode/language_server/common/protocol.py

@@ -41,6 +41,7 @@
 from .parts.declaration import DeclarationProtocolPart
 from .parts.definition import DefinitionProtocolPart
 from .parts.diagnostics import DiagnosticsProtocolPart
+from .parts.document_highlight import DocumentHighlightProtocolPart
 from .parts.document_symbols import DocumentSymbolsProtocolPart
 from .parts.documents import TextDocumentProtocolPart
 from .parts.folding_range import FoldingRangeProtocolPart
@@ -79,6 +80,7 @@ class LanguageServerProtocol(JsonRPCProtocol):
     formatting = ProtocolPartDescriptor(FormattingProtocolPart)
     semantic_tokens = ProtocolPartDescriptor(SemanticTokensProtocolPart)
     references = ProtocolPartDescriptor(ReferencesProtocolPart)
+    document_highlight = ProtocolPartDescriptor(DocumentHighlightProtocolPart)
 
     name: Optional[str] = None
     version: Optional[str] = None

robotcode/language_server/robotframework/diagnostics/library_doc.py

@@ -165,20 +165,11 @@ class InvalidVariableError(Exception):
 
 
 class VariableMatcher:
-    _match_extended = re.compile(
-        r"""
-        (.+?)          # base name (group 1)
-        ([^\s\w].+)    # extended part (group 2)
-        """,
-        re.UNICODE | re.VERBOSE,
-    )
-
-    def __init__(self, name: str, extended: bool = False) -> None:
+    def __init__(self, name: str) -> None:
         from robot.utils.normalizing import normalize
         from robot.variables.search import VariableSearcher
 
         self.name = name
-        self.extended = extended
 
         searcher = VariableSearcher("$@&%", ignore_errors=True)
         match = searcher.search(name)
@@ -188,28 +179,20 @@ def __init__(self, name: str, extended: bool = False) -> None:
 
         self.base = match.base
 
-        self.extended_base: Optional[str] = None
-        self.normalize_extended: Optional[str] = None
-        if extended:
-            ext_match = self._match_extended.match(self.name[2:-1])
-            if ext_match is not None:
-                self.extended_base, _ = ext_match.groups()
-                self.normalize_extended = str(normalize(self.extended_base, "_"))
-
         self.normalized_name = str(normalize(self.base, "_"))
 
     def __eq__(self, o: object) -> bool:
         from robot.utils.normalizing import normalize
         from robot.variables.search import VariableSearcher
 
         if isinstance(o, VariableMatcher):
-            return o.normalized_name == self.normalized_name or (o.extended and self.normalized_name == o.extended_base)
+            return o.normalized_name == self.normalized_name
         elif isinstance(o, str):
             searcher = VariableSearcher("$@&%", ignore_errors=True)
             match = searcher.search(o)
             base = match.base
             normalized = str(normalize(base, "_"))
-            return self.normalized_name == normalized or (self.extended and self.normalize_extended == normalized)
+            return self.normalized_name == normalized
         else:
             return False
 

robotcode/language_server/robotframework/diagnostics/namespace.py

@@ -3,6 +3,7 @@
 import ast
 import asyncio
 import itertools
+import re
 import weakref
 from collections import OrderedDict
 from dataclasses import dataclass, field
@@ -588,14 +589,12 @@ def get_command_line_variables(self) -> List[VariableDefinition]:
         return self.imports_manager.get_command_line_variables()
 
     @_logger.call
-    async def get_variables(
+    async def yield_variables(
         self, nodes: Optional[List[ast.AST]] = None, position: Optional[Position] = None
-    ) -> Dict[VariableMatcher, VariableDefinition]:
+    ) -> AsyncGenerator[Tuple[VariableMatcher, VariableDefinition], None]:
         from robot.parsing.model.blocks import Keyword, TestCase
 
-        await self.ensure_initialized()
-
-        result: Dict[VariableMatcher, VariableDefinition] = {}
+        yielded: Dict[VariableMatcher, VariableDefinition] = {}
 
         async for var in async_chain(
             *[
@@ -609,16 +608,46 @@ async def get_variables(
             *(e.variables for e in self._variables.values()),
             (e for e in self.get_builtin_variables()),
         ):
-            if var.name is not None and VariableMatcher(var.name) not in result.keys():
-                result[VariableMatcher(var.name)] = var
+            if var.name is not None:
+                matcher = VariableMatcher(var.name)
+                if matcher not in yielded.keys():
+                    yielded[matcher] = var
+                    yield matcher, var
 
-        return result
+    @_logger.call
+    async def get_variables(
+        self, nodes: Optional[List[ast.AST]] = None, position: Optional[Position] = None
+    ) -> Dict[VariableMatcher, VariableDefinition]:
+        return {m: v async for m, v in self.yield_variables(nodes, position)}
+
+    _match_extended = re.compile(
+        r"""
+        (.+?)          # base name (group 1)
+        ([^\s\w].+)    # extended part (group 2)
+        """,
+        re.UNICODE | re.VERBOSE,
+    )
 
     @_logger.call
     async def find_variable(
         self, name: str, nodes: Optional[List[ast.AST]], position: Optional[Position] = None
     ) -> Optional[VariableDefinition]:
-        return (await self.get_variables(nodes, position)).get(VariableMatcher(name, True), None)
+
+        matcher = VariableMatcher(name)
+        async for m, v in self.yield_variables(nodes, position):
+            if matcher == m:
+                return v
+
+        match = self._match_extended.match(name[2:-1])
+        if match is not None:
+            base_name, extended = match.groups()
+            name = f"{name[0]}{{{base_name}}}"
+
+            matcher = VariableMatcher(name)
+            async for m, v in self.yield_variables(nodes, position):
+                if matcher == m:
+                    return v
+        return None
 
     @_logger.call
     async def _import_imports(self, imports: Iterable[Import], base_dir: str, *, top_level: bool = False) -> None: