implement find references for variables

Author: d-biehl

CHANGELOG.md

@@ -9,6 +9,7 @@ All notable changes to the "robotcode" extension will be documented in this file
   - goto
   - static and dynamic variables
 - correct debugger hover on variables and last fail message
+- implement find references for variables
 
 
 ##  0.4.1

robotcode/language_server/robotframework/diagnostics/entities.py

@@ -76,6 +76,7 @@ def __hash__(self) -> int:
 
 class VariableDefinitionType(Enum):
     VARIABLE = "variable"
+    LOCAL_VARIABLE = "local variable"
     ARGUMENT = "argument"
     COMMAND_LINE_VARIABLE = "command line variable"
     BUILTIN_VARIABLE = "builtin variable"
@@ -104,6 +105,14 @@ def range(self) -> Range:
         )
 
 
+@dataclass
+class LocalVariableDefinition(VariableDefinition):
+    type: VariableDefinitionType = VariableDefinitionType.LOCAL_VARIABLE
+
+    def __hash__(self) -> int:
+        return hash((type(self), self.name, self.type))
+
+
 @dataclass
 class BuiltInVariableDefinition(VariableDefinition):
     type: VariableDefinitionType = VariableDefinitionType.BUILTIN_VARIABLE

robotcode/language_server/robotframework/diagnostics/imports_manager.py

@@ -29,6 +29,7 @@
 from ...common.text_document import TextDocument
 from ..configuration import RobotConfig
 from ..utils.async_ast import walk
+from .entities import CommandLineVariableDefinition, VariableDefinition
 
 if TYPE_CHECKING:
     from ..protocol import RobotLanguageServerProtocol
@@ -465,6 +466,20 @@ def __init__(
         self._variables: OrderedDict[_VariablesEntryKey, _VariablesEntry] = OrderedDict()
         self.file_watchers: List[FileWatcherEntry] = []
         self.parent_protocol.documents.did_change.add(self.resource_document_changed)
+        self._command_line_variables: Optional[List[VariableDefinition]] = None
+
+    @_logger.call
+    def get_command_line_variables(self) -> List[VariableDefinition]:
+        if self._command_line_variables is None:
+            if self.config is None:
+                self._command_line_variables = []
+            else:
+                self._command_line_variables = [
+                    CommandLineVariableDefinition(0, 0, 0, 0, "", f"${{{k}}}", None)
+                    for k in self.config.variables.keys()
+                ]
+
+        return self._command_line_variables
 
     @async_tasking_event
     async def libraries_changed(sender, libraries: List[LibraryDoc]) -> None:  # NOSONAR

robotcode/language_server/robotframework/diagnostics/library_doc.py

@@ -165,7 +165,15 @@ class InvalidVariableError(Exception):
 
 
 class VariableMatcher:
-    def __init__(self, name: str) -> None:
+    _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:
         from robot.utils.normalizing import normalize
         from robot.variables.search import VariableSearcher
 
@@ -176,7 +184,13 @@ def __init__(self, name: str) -> None:
             raise InvalidVariableError(f"Invalid variable '{name}'")
 
         self.base = match.base
-        self.normalized_name = str(normalize(match.base, "_"))
+
+        if extended:
+            ext_match = self._match_extended.match(self.name[2:-1])
+            if ext_match is not None:
+                self.base, _ = ext_match.groups()
+
+        self.normalized_name = str(normalize(self.base, "_"))
 
     def __eq__(self, o: object) -> bool:
         from robot.utils.normalizing import normalize
@@ -1224,21 +1238,18 @@ def get_variables_doc(
     stem = Path(name).stem
     try:
         source = find_file(name, working_dir, base_dir, pythonpath, environment, variables)
+        with _std_capture() as std_capturer:
 
-        if source.lower().endswith((".yaml", ".yml")):
-            importer = YamlImporter()
-        else:
-            importer = PythonImporter()
-        vars: List[VariableDefinition] = [
-            ImportedVariableDefinition(1, 0, 1, 0, source, var[0], None)
-            for var in importer.import_variables(source, args)
-        ]
+            if source.lower().endswith((".yaml", ".yml")):
+                importer = YamlImporter()
+            else:
+                importer = PythonImporter()
+            vars: List[VariableDefinition] = [
+                ImportedVariableDefinition(1, 0, 1, 0, source, var[0], None)
+                for var in importer.import_variables(source, args)
+            ]
 
-        return VariablesDoc(
-            name=stem,
-            source=source,
-            variables=vars,
-        )
+            return VariablesDoc(name=stem, source=source, variables=vars, stdout=std_capturer.getvalue())
     except BaseException as e:
         return VariablesDoc(
             name=stem,

robotcode/language_server/robotframework/diagnostics/namespace.py

@@ -39,9 +39,9 @@
 from .entities import (
     ArgumentDefinition,
     BuiltInVariableDefinition,
-    CommandLineVariableDefinition,
     Import,
     LibraryImport,
+    LocalVariableDefinition,
     ResourceImport,
     VariableDefinition,
     VariablesImport,
@@ -206,7 +206,7 @@ async def visit_KeywordCall(self, node: ast.AST) -> None:  # noqa: N802
             variable_token = self.get_variable_token(assign_token)
             try:
                 if variable_token is not None and variable_token.value not in self._results:
-                    self._results[variable_token.value] = VariableDefinition(
+                    self._results[variable_token.value] = LocalVariableDefinition(
                         name=variable_token.value,
                         name_token=variable_token,
                         line_no=variable_token.lineno,
@@ -228,7 +228,7 @@ async def visit_ForHeader(self, node: ast.AST) -> None:  # noqa: N802
         for variable in variables:
             variable_token = self.get_variable_token(variable)
             if variable_token is not None and variable_token.value and variable_token.value not in self._results:
-                self._results[variable_token.value] = VariableDefinition(
+                self._results[variable_token.value] = LocalVariableDefinition(
                     name=variable_token.value,
                     name_token=variable_token,
                     line_no=node.lineno,
@@ -580,13 +580,7 @@ def get_builtin_variables(cls) -> List[BuiltInVariableDefinition]:
 
     @_logger.call
     def get_command_line_variables(self) -> List[VariableDefinition]:
-        if self.imports_manager.config is None:
-            return []
-
-        return [
-            CommandLineVariableDefinition(0, 0, 0, 0, "", f"${{{k}}}", None)
-            for k in self.imports_manager.config.variables.keys()
-        ]
+        return self.imports_manager.get_command_line_variables()
 
     @_logger.call
     async def get_variables(
@@ -619,7 +613,7 @@ async def get_variables(
     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), None)
+        return (await self.get_variables(nodes, position)).get(VariableMatcher(name, True), None)
 
     @_logger.call
     async def _import_imports(self, imports: Iterable[Import], base_dir: str, *, top_level: bool = False) -> None:

robotcode/language_server/robotframework/parts/completion.py

@@ -338,10 +338,12 @@ async def create_environment_variables_completion_items(self, range: Optional[Ra
         ]
 
     _VARIABLE_COMPLETION_SORT_TEXT_PREFIX = {
-        VariableDefinitionType.VARIABLE: "035",
+        VariableDefinitionType.LOCAL_VARIABLE: "033",
         VariableDefinitionType.ARGUMENT: "034",
-        VariableDefinitionType.COMMAND_LINE_VARIABLE: "036",
-        VariableDefinitionType.BUILTIN_VARIABLE: "037",
+        VariableDefinitionType.VARIABLE: "035",
+        VariableDefinitionType.IMPORTED_VARIABLE: "036",
+        VariableDefinitionType.COMMAND_LINE_VARIABLE: "037",
+        VariableDefinitionType.BUILTIN_VARIABLE: "038",
     }
 
     async def create_variables_completion_items(
@@ -532,7 +534,6 @@ def enumerate_indexes(s: str, c: str) -> Iterator[int]:
                 detail="Library",
                 sort_text=f"030_{k}",
                 deprecated=v.library_doc.is_deprecated,
-                # documentation=MarkupContent(kind=MarkupKind.MARKDOWN, value=v.library_doc.to_markdown()),
                 insert_text_format=InsertTextFormat.PLAINTEXT,
                 text_edit=TextEdit(range=r, new_text=k) if r is not None else None,
                 data={
@@ -550,7 +551,6 @@ def enumerate_indexes(s: str, c: str) -> Iterator[int]:
                 detail="Resource",
                 deprecated=v.library_doc.is_deprecated,
                 sort_text=f"030_{k}",
-                # documentation=MarkupContent(kind=MarkupKind.MARKDOWN, value=v.library_doc.to_markdown()),
                 insert_text_format=InsertTextFormat.PLAINTEXT,
                 text_edit=TextEdit(range=r, new_text=k) if r is not None else None,
                 data={

robotcode/language_server/robotframework/parts/references.py

@@ -17,7 +17,6 @@
 )
 
 from ....utils.async_tools import (
-    CancelationToken,
     check_canceled,
     create_sub_task,
     run_coroutine_in_thread,
@@ -29,6 +28,11 @@
 from ...common.lsp_types import Location, Position, ReferenceContext
 from ...common.text_document import TextDocument
 from ..configuration import WorkspaceConfig
+from ..diagnostics.entities import (
+    ArgumentDefinition,
+    LocalVariableDefinition,
+    VariableDefinition,
+)
 from ..diagnostics.library_doc import (
     ALL_RUN_KEYWORDS_MATCHERS,
     RESOURCE_FILE_EXTENSION,
@@ -136,12 +140,19 @@ async def _references_default(
 
                     if position.is_in_range(range):
                         variable = await namespace.find_variable(sub_token.value, nodes, position)
-                        if variable is not None and variable.source:
+                        if variable is not None:
                             return [
-                                Location(
-                                    uri=str(Uri.from_path(variable.source)),
-                                    range=variable.range(),
-                                )
+                                *(
+                                    [
+                                        Location(
+                                            uri=str(Uri.from_path(variable.source)),
+                                            range=variable.range(),
+                                        ),
+                                    ]
+                                    if context.include_declaration and variable.source
+                                    else []
+                                ),
+                                *await self._find_variable_references(document, variable),
                             ]
             except (SystemExit, KeyboardInterrupt, asyncio.CancelledError):
                 raise
@@ -298,11 +309,10 @@ async def _find_keyword_references_in_file(
         kw_doc: KeywordDoc,
         lib_doc: Optional[LibraryDoc],
         file: Path,
-        cancel_token: CancelationToken,
     ) -> List[Location]:
 
         doc = await self.parent.robot_workspace.get_or_open_document(file, "robotframework")
-        namespace = await self.parent.documents_cache.get_namespace(doc, cancelation_token=cancel_token)
+        namespace = await self.parent.documents_cache.get_namespace(doc)
 
         if (
             lib_doc is not None
@@ -313,11 +323,9 @@ async def _find_keyword_references_in_file(
         ):
             return []
 
-        return await self._find_keyword_references_in_namespace(namespace, kw_doc, cancel_token)
+        return await self._find_keyword_references_in_namespace(namespace, kw_doc)
 
-    async def _find_keyword_references_in_namespace(
-        self, namespace: Namespace, kw_doc: KeywordDoc, cancel_token: CancelationToken
-    ) -> List[Location]:
+    async def _find_keyword_references_in_namespace(self, namespace: Namespace, kw_doc: KeywordDoc) -> List[Location]:
         from robot.parsing.lexer.tokens import Token as RobotToken
         from robot.parsing.model.statements import (
             Fixture,
@@ -503,8 +511,6 @@ async def _find_keyword_references(self, document: TextDocument, kw_doc: Keyword
             or await namespace.get_library_doc()
         )
 
-        cancel_token = CancelationToken()
-
         futures: List[Awaitable[List[Location]]] = []
         result: List[Location] = []
 
@@ -516,7 +522,77 @@ async def _find_keyword_references(self, document: TextDocument, kw_doc: Keyword
             ignore_patterns=config.exclude_patterns or [],  # type: ignore
             absolute=True,
         ):
-            futures.append(create_sub_task(self._find_keyword_references_in_file(kw_doc, lib_doc, f, cancel_token)))
+            futures.append(create_sub_task(self._find_keyword_references_in_file(kw_doc, lib_doc, f)))
+
+        for e in await asyncio.gather(*futures, return_exceptions=True):
+            if isinstance(e, BaseException):
+                self._logger.exception(e)
+                continue
+            result.extend(e)
+
+        return result
+
+    @_logger.call
+    async def _find_variable_references_in_file(
+        self,
+        variable: VariableDefinition,
+        file: Path,
+    ) -> List[Location]:
+        from robot.parsing.lexer.tokens import Token as RobotToken
+        from robot.parsing.model.blocks import Block, Keyword, Section, TestCase
+
+        doc = await self.parent.robot_workspace.get_or_open_document(file, "robotframework")
+        namespace = await self.parent.documents_cache.get_namespace(doc)
+        model = await self.parent.documents_cache.get_model(doc)
+
+        result: List[Location] = []
+        current_block: Optional[Block] = None
+
+        async for node in iter_nodes(model):
+            if isinstance(node, Section):
+                current_block = None
+            elif isinstance(node, (TestCase, Keyword)):
+                current_block = node
+
+            if isinstance(node, HasTokens):
+                for token in node.tokens:
+                    for sub_token in tokenize_variables(token):
+                        if sub_token.type == RobotToken.VARIABLE:
+                            found_variable = await namespace.find_variable(
+                                sub_token.value,
+                                [*([current_block] if current_block is not None else []), node],
+                                range_from_token(token).start,
+                            )
+
+                            if found_variable == variable:
+                                result.append(Location(str(doc.uri), range_from_token(sub_token)))
+
+        return result
+
+    async def _find_variable_references(self, document: TextDocument, variable: VariableDefinition) -> List[Location]:
+        folder = self.parent.workspace.get_workspace_folder(document.uri)
+        if folder is None:
+            return []
+
+        namespace = await self.parent.documents_cache.get_namespace(document)
+        if namespace is None:
+            return None
+
+        futures: List[Awaitable[List[Location]]] = []
+        result: List[Location] = []
+
+        config = await self.parent.workspace.get_configuration(WorkspaceConfig, folder.uri) or WorkspaceConfig()
+
+        if isinstance(variable, (ArgumentDefinition, LocalVariableDefinition)):
+            futures.append(create_sub_task(self._find_variable_references_in_file(variable, document.uri.to_path())))
+        else:
+            async for f in iter_files(
+                folder.uri.to_path(),
+                (f"**/*.{{{ROBOT_FILE_EXTENSION[1:]},{RESOURCE_FILE_EXTENSION[1:]}}}"),
+                ignore_patterns=config.exclude_patterns or [],  # type: ignore
+                absolute=True,
+            ):
+                futures.append(create_sub_task(self._find_variable_references_in_file(variable, f)))
 
         for e in await asyncio.gather(*futures, return_exceptions=True):
             if isinstance(e, BaseException):