correct handling of argument definition wich contains a default value with variable

Author: d-biehl

robotcode/language_server/robotframework/diagnostics/namespace.py

@@ -141,12 +141,19 @@ async def visit_Variable(self, node: ast.AST) -> None:  # noqa: N802
 
 
 class BlockVariableVisitor(AsyncVisitor):
-    async def get(self, source: str, model: ast.AST, position: Optional[Position] = None) -> List[VariableDefinition]:
+    def __init__(self, source: str, position: Optional[Position] = None, in_args: bool = True) -> None:
+        super().__init__()
+
         self.source = source
         self.position = position
+        self.in_args = in_args
 
         self._results: Dict[str, VariableDefinition] = {}
 
+    async def get(self, model: ast.AST) -> List[VariableDefinition]:
+
+        self._results = {}
+
         await self.visit(model)
 
         return list(self._results.values())
@@ -178,10 +185,10 @@ async def visit_KeywordName(self, node: ast.AST) -> None:  # noqa: N802
                     self._results[name] = ArgumentDefinition(
                         name=name,
                         name_token=strip_variable_token(variable_token),
-                        line_no=n.lineno,
-                        col_offset=n.col_offset,
-                        end_line_no=n.lineno,
-                        end_col_offset=n.end_col_offset,
+                        line_no=variable_token.lineno,
+                        col_offset=variable_token.col_offset,
+                        end_line_no=variable_token.lineno,
+                        end_col_offset=variable_token.end_col_offset,
                         source=self.source,
                     )
 
@@ -212,13 +219,16 @@ async def visit_Arguments(self, node: ast.AST) -> None:  # noqa: N802
                 argument = self.get_variable_token(argument_token)
 
                 if argument is not None:
+                    if self.in_args and self.position is not None and self.position > range_from_token(argument).end:
+                        break
+
                     self._results[argument.value] = ArgumentDefinition(
                         name=argument.value,
                         name_token=strip_variable_token(argument),
-                        line_no=n.lineno,
-                        col_offset=n.col_offset,
-                        end_line_no=n.lineno,
-                        end_col_offset=n.end_col_offset,
+                        line_no=argument.lineno,
+                        col_offset=argument.col_offset,
+                        end_line_no=argument.lineno,
+                        end_col_offset=argument.end_col_offset,
                         source=self.source,
                     )
 
@@ -241,10 +251,10 @@ async def visit_ExceptHeader(self, node: ast.AST) -> None:  # noqa: N802
                     self._results[variable.value] = LocalVariableDefinition(
                         name=variable.value,
                         name_token=strip_variable_token(variable),
-                        line_no=n.lineno,
-                        col_offset=n.col_offset,
-                        end_line_no=n.lineno,
-                        end_col_offset=n.end_col_offset,
+                        line_no=variable.lineno,
+                        col_offset=variable.col_offset,
+                        end_line_no=variable.lineno,
+                        end_col_offset=variable.end_col_offset,
                         source=self.source,
                     )
 
@@ -276,10 +286,10 @@ async def visit_KeywordCall(self, node: ast.AST) -> None:  # noqa: N802
                         self._results[variable_token.value] = LocalVariableDefinition(
                             name=variable_token.value,
                             name_token=strip_variable_token(variable_token),
-                            line_no=n.lineno,
-                            col_offset=n.col_offset,
-                            end_line_no=n.lineno,
-                            end_col_offset=n.end_col_offset,
+                            line_no=variable_token.lineno,
+                            col_offset=variable_token.col_offset,
+                            end_line_no=variable_token.lineno,
+                            end_col_offset=variable_token.end_col_offset,
                             source=self.source,
                         )
 
@@ -311,10 +321,10 @@ async def visit_InlineIfHeader(self, node: ast.AST) -> None:  # noqa: N802
                         self._results[variable_token.value] = LocalVariableDefinition(
                             name=variable_token.value,
                             name_token=strip_variable_token(variable_token),
-                            line_no=n.lineno,
-                            col_offset=n.col_offset,
-                            end_line_no=n.lineno,
-                            end_col_offset=n.end_col_offset,
+                            line_no=variable_token.lineno,
+                            col_offset=variable_token.col_offset,
+                            end_line_no=variable_token.lineno,
+                            end_col_offset=variable_token.end_col_offset,
                             source=self.source,
                         )
 
@@ -333,10 +343,10 @@ async def visit_ForHeader(self, node: ast.AST) -> None:  # noqa: N802
                 self._results[variable_token.value] = LocalVariableDefinition(
                     name=variable_token.value,
                     name_token=strip_variable_token(variable_token),
-                    line_no=n.lineno,
-                    col_offset=n.col_offset,
-                    end_line_no=n.lineno,
-                    end_col_offset=n.end_col_offset,
+                    line_no=variable_token.lineno,
+                    col_offset=variable_token.col_offset,
+                    end_line_no=variable_token.lineno,
+                    end_col_offset=variable_token.end_col_offset,
                     source=self.source,
                 )
 
@@ -749,14 +759,17 @@ async def yield_variables(
         skip_commandline_variables: bool = False,
     ) -> AsyncGenerator[Tuple[VariableMatcher, VariableDefinition], None]:
         from robot.parsing.model.blocks import Keyword, TestCase
+        from robot.parsing.model.statements import Arguments
 
         # await self.ensure_initialized()
 
         yielded: Dict[VariableMatcher, VariableDefinition] = {}
 
         async for var in async_chain(
             *[
-                await BlockVariableVisitor().get(self.source, n, position)
+                await BlockVariableVisitor(
+                    self.source, position, isinstance(nodes[-1], Arguments) if nodes else False
+                ).get(n)
                 for n in nodes or []
                 if isinstance(n, (Keyword, TestCase))
             ],

robotcode/language_server/robotframework/parts/references.py

@@ -2,6 +2,7 @@
 
 import ast
 import asyncio
+import itertools
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -47,8 +48,10 @@
     get_nodes_at_position,
     get_tokens_at_position,
     is_not_variable_token,
+    iter_nodes_at_position,
     iter_over_keyword_names_and_owners,
     range_from_token,
+    tokenize_variables,
 )
 from ..utils.async_ast import iter_nodes
 
@@ -227,21 +230,46 @@ async def find_variable_references(self, document: TextDocument, variable: Varia
             else await self._find_references(document, self.find_variable_references_in_file, variable)
         )
 
+    async def yield_argument_name_and_rest(self, node: ast.AST, token: Token) -> AsyncGenerator[Token, None]:
+        from robot.parsing.lexer.tokens import Token as RobotToken
+        from robot.parsing.model.statements import Arguments
+
+        if isinstance(node, Arguments) and token.type == RobotToken.ARGUMENT:
+            argument = next(
+                (
+                    v
+                    for v in itertools.dropwhile(
+                        lambda t: t.type in RobotToken.NON_DATA_TOKENS,
+                        tokenize_variables(token, ignore_errors=True),
+                    )
+                    if v.type == RobotToken.VARIABLE
+                ),
+                None,
+            )
+            if argument is None or argument.value == token.value:
+                yield token
+            else:
+                yield argument
+                i = len(argument.value)
+
+                async for t in self.yield_argument_name_and_rest(
+                    node, RobotToken(token.type, token.value[i:], token.lineno, token.col_offset + i, token.error)
+                ):
+                    yield t
+        else:
+            yield token
+
     @_logger.call
     async def find_variable_references_in_file(
         self,
         doc: TextDocument,
         variable: VariableDefinition,
     ) -> List[Location]:
         from robot.parsing.lexer.tokens import Token as RobotToken
-        from robot.parsing.model.blocks import Block, Keyword, Section, TestCase
 
         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
-
         if (
             variable.source
             and variable.source != str(doc.uri.to_path())
@@ -259,24 +287,23 @@ async def find_variable_references_in_file(
 
         expression_statements = self.get_expression_statement_types()
 
+        result: List[Location] = []
+
         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:
-                    async for token_and_var in self.iter_variables_from_token(
-                        token,
-                        namespace,
-                        [*([current_block] if current_block is not None else []), node],
-                        range_from_token(token).start,
-                    ):
-                        sub_token, found_variable = token_and_var
-
-                        if found_variable == variable:
-                            result.append(Location(str(doc.uri), range_from_token(sub_token)))
+                for token1 in node.tokens:
+                    async for token in self.yield_argument_name_and_rest(node, token1):
+                        async for token_and_var in self.iter_variables_from_token(
+                            token,
+                            namespace,
+                            [n async for n in iter_nodes_at_position(model, range_from_token(token).start)],
+                            range_from_token(token).start,
+                        ):
+                            sub_token, found_variable = token_and_var
+
+                            if found_variable == variable:
+                                result.append(Location(str(doc.uri), range_from_token(sub_token)))
 
             if (
                 isinstance(node, Statement)
@@ -286,7 +313,7 @@ async def find_variable_references_in_file(
                 async for token_and_var in self.iter_expression_variables_from_token(
                     token,
                     namespace,
-                    [*([current_block] if current_block is not None else []), node],
+                    [n async for n in iter_nodes_at_position(model, range_from_token(token).start)],
                     range_from_token(token).start,
                 ):
                     sub_token, found_variable = token_and_var