implement hover and goto for "FOR" variables

Author: d-biehl

robotcode/language_server/robotframework/diagnostics/namespace.py

@@ -38,6 +38,7 @@
 from ..utils.ast import (
     Token,
     is_non_variable_token,
+    range_from_node,
     range_from_token_or_node,
     tokenize_variables,
 )
@@ -219,12 +220,20 @@ async def visit_Variable(self, node: ast.AST) -> None:  # noqa: N802
 
 
 class BlockVariableVisitor(AsyncVisitor):
-    async def get(self, source: str, model: ast.AST) -> List[VariableDefinition]:
-        self._results: List[VariableDefinition] = []
+    async def get(self, source: str, model: ast.AST, position: Optional[Position] = None) -> List[VariableDefinition]:
         self.source = source
+        self.position = position
+
+        self._results: List[VariableDefinition] = []
+
         await self.visit(model)
+
         return self._results
 
+    async def visit(self, node: ast.AST) -> None:
+        if self.position is None or self.position >= range_from_node(node).start:
+            return await super().visit(node)
+
     async def visit_KeywordName(self, node: ast.AST) -> None:  # noqa: N802
         from robot.parsing.lexer.tokens import Token as RobotToken
         from robot.parsing.model.statements import KeywordName
@@ -343,6 +352,39 @@ async def visit_KeywordCall(self, node: ast.AST) -> None:  # noqa: N802
         except VariableError:
             pass
 
+    async def visit_ForHeader(self, node: ast.AST) -> None:  # noqa: N802
+        from robot.errors import VariableError
+        from robot.parsing.lexer.tokens import Token as RobotToken
+        from robot.parsing.model.statements import ForHeader
+        from robot.variables.search import contains_variable
+
+        try:
+            n = cast(ForHeader, node)
+            variables = n.get_tokens(RobotToken.VARIABLE)
+            for variable in variables:
+                if variable is not None and variable.value and contains_variable(variable.value):
+                    self._results.append(
+                        VariableDefinition(
+                            name=variable.value,
+                            name_token=variable,
+                            line_no=node.lineno,
+                            col_offset=node.col_offset,
+                            end_line_no=node.end_lineno
+                            if node.end_lineno is not None
+                            else variable.lineno
+                            if variable.lineno is not None
+                            else -1,
+                            end_col_offset=node.end_col_offset
+                            if node.end_col_offset is not None
+                            else variable.end_col_offset
+                            if variable.end_col_offset is not None
+                            else -1,
+                            source=self.source,
+                        )
+                    )
+        except VariableError:
+            pass
+
 
 class ImportVisitor(AsyncVisitor):
     async def get(self, source: str, model: ast.AST) -> List[Import]:
@@ -953,7 +995,9 @@ def get_builtin_variables(cls) -> List[BuiltInVariableDefinition]:
 
         return cls._builtin_variables
 
-    async def get_variables(self, nodes: Optional[List[ast.AST]] = None) -> Dict[VariableMatcher, VariableDefinition]:
+    async def get_variables(
+        self, nodes: Optional[List[ast.AST]] = None, position: Optional[Position] = None
+    ) -> Dict[VariableMatcher, VariableDefinition]:
         from robot.parsing.model.blocks import Keyword, TestCase
 
         await self.ensure_initialized()
@@ -962,7 +1006,7 @@ async def get_variables(self, nodes: Optional[List[ast.AST]] = None) -> Dict[Var
 
         async for var in async_chain(
             *[
-                await BlockVariableVisitor().get(self.source, n)
+                await BlockVariableVisitor().get(self.source, n, position)
                 for n in nodes or []
                 if isinstance(n, (Keyword, TestCase))
             ],
@@ -975,8 +1019,10 @@ async def get_variables(self, nodes: Optional[List[ast.AST]] = None) -> Dict[Var
 
         return result
 
-    async def find_variable(self, name: str, nodes: Optional[List[ast.AST]]) -> Optional[VariableDefinition]:
-        return (await self.get_variables(nodes)).get(VariableMatcher(name), None)
+    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)
 
     async def _import_imports(self, imports: Iterable[Import], base_dir: str, *, top_level: bool = False) -> None:
         async def _import(value: Import) -> Optional[LibraryEntry]:

robotcode/language_server/robotframework/parts/completion.py

@@ -330,7 +330,7 @@ async def create_variables_completion_items(
                 ),
                 filter_text=s.name[2:-1] if range is not None else None,
             )
-            for s in (await namespace.get_variables(nodes)).values()
+            for s in (await namespace.get_variables(nodes, position)).values()
             if s.name is not None and (s.name_token is None or not position.is_in_range(range_from_token(s.name_token)))
         ]
 

robotcode/language_server/robotframework/parts/goto.py

@@ -112,7 +112,7 @@ async def _definition_default(
                     range = range_from_token(sub_token)
 
                     if position.is_in_range(range):
-                        variable = await namespace.find_variable(sub_token.value, nodes)
+                        variable = await namespace.find_variable(sub_token.value, nodes, position)
                         if variable is not None and variable.source:
                             return [
                                 LocationLink(

robotcode/language_server/robotframework/parts/hover.py

@@ -100,7 +100,7 @@ async def _hover_default(self, nodes: List[ast.AST], document: TextDocument, pos
                     range = range_from_token(sub_token)
 
                     if position.is_in_range(range):
-                        variable = await namespace.find_variable(sub_token.value, nodes)
+                        variable = await namespace.find_variable(sub_token.value, nodes, position)
                         if variable is not None:
                             return Hover(
                                 contents=MarkupContent(

tests/robotcode/language_server/robotframework/parts/test_hover.py

@@ -49,7 +49,11 @@ async def test_document() -> AsyncGenerator[TextDocument, None]:
     data_path = Path(Path(__file__).parent, "data/hover.robot")
     data = data_path.read_text()
 
-    yield TextDocument(document_uri=data_path.as_uri(), language_id="robotframework", version=1, text=data)
+    document = TextDocument(document_uri=data_path.as_uri(), language_id="robotframework", version=1, text=data)
+    try:
+        yield document
+    finally:
+        del document
 
 
 @pytest.mark.parametrize(
@@ -87,3 +91,25 @@ async def test_hover_should_not_find_simple_keyword_on_boundaries(
 
     result = await protocol._robot_hover.collect(protocol.hover, test_document, position)
     assert result is None
+
+
+@pytest.mark.parametrize(
+    ("position", "variable"),
+    [
+        (Position(line=4, character=2), "(Variable) ${A VAR}"),
+        (Position(line=9, character=18), "(Variable) ${A VAR}"),
+        (Position(line=5, character=7), "(Variable) &{A DICT}"),
+        (Position(line=10, character=36), "(Variable) &{A DICT}"),
+        (Position(line=11, character=13), "(Variable) ${key}"),  # FOR Variable
+        (Position(line=11, character=24), "(Variable) ${value}"),  # FOR Variable
+    ],
+)
+@pytest.mark.asyncio
+async def test_hover_should_find_variable(
+    protocol: RobotLanguageServerProtocol, test_document: TextDocument, position: Position, variable: str
+) -> None:
+
+    result = await protocol._robot_hover.collect(protocol.hover, test_document, position)
+    assert result
+    assert isinstance(result.contents, MarkupContent)
+    assert result.contents.value == variable