find declaration of local variables

vyper_lsp/ast.py

@@ -270,6 +270,16 @@ def find_nodes_referencing_symbol(self, symbol: str):
 
         return return_nodes
 
+    def find_node_declaring_symbol(self, symbol: str):
+        if self.ast_data is None:
+            return None
+
+        for node in self.ast_data.get_descendants(
+            (nodes.AnnAssign, nodes.VariableDecl)
+        ):
+            if node.target.id == symbol:
+                return node
+
     @classmethod
     def create_new_instance(cls, ast):
         # Create a new instance

vyper_lsp/navigation.py

@@ -1,10 +1,9 @@
 import re
-import sys
 from pygls.lsp.types.language_features import Position, Range
 from typing import List, Optional
 
 from pygls.workspace import Document
-from vyper.ast import EnumDef, FunctionDef
+from vyper.ast import EnumDef, FunctionDef, VyperNode
 from vyper_lsp.ast import AST
 from vyper_lsp.utils import get_expression_at_cursor, get_word_at_cursor
 
@@ -26,6 +25,21 @@ def find_state_variable_declaration(self, word: str) -> Optional[Range]:
             )
             return range
 
+    def find_variable_declaration_under_node(
+        self, node: VyperNode, symbol: str
+    ) -> Optional[Range]:
+        decl_node = AST.create_new_instance(node).find_node_declaring_symbol(symbol)
+        if decl_node:
+            range = Range(
+                start=Position(
+                    line=decl_node.lineno - 1, character=decl_node.col_offset
+                ),
+                end=Position(
+                    line=decl_node.end_lineno - 1, character=decl_node.end_col_offset
+                ),
+            )
+            return range
+
     def find_function_declaration(self, word: str) -> Optional[Range]:
         node = self.ast.find_function_declaration_node_for_name(word)
         if node:
@@ -123,8 +137,7 @@ def find_declaration(self, document: Document, pos: Position) -> Optional[Range]
         word = get_word_at_cursor(og_line, pos.character)
         full_word = get_expression_at_cursor(og_line, pos.character)
         range = None
-
-        print(f"word: {word}", file=sys.stderr)
+        top_level_node = self.ast.find_top_level_node_at_pos(pos)
 
         if full_word.startswith("self."):
             if "(" in full_word:
@@ -136,17 +149,19 @@ def find_declaration(self, document: Document, pos: Position) -> Optional[Range]
                 range = self.find_type_declaration(word)
             elif word in self.ast.get_constants():
                 range = self.find_state_variable_declaration(word)
-            else:
-                # check if full_word matches "enum.variant" regex
-                pattern = r"([a-zA-Z_][a-zA-Z0-9_]*)\.([a-zA-Z_][a-zA-Z0-9_]*)"
-                match = re.match(pattern, full_word)
-                if match:
-                    enum_name = match.group(1)
-                    variant_name = match.group(2)
-                    if enum_name in self.ast.get_enums():
-                        variants = self.ast.get_enum_variants(enum_name)
-                        if variant_name in variants:
-                            range = self.find_type_declaration(enum_name)
+            elif isinstance(top_level_node, FunctionDef):
+                range = self.find_variable_declaration_under_node(top_level_node, word)
+                if not range:
+                    # check if full_word matches "enum.variant" regex
+                    pattern = r"([a-zA-Z_][a-zA-Z0-9_]*)\.([a-zA-Z_][a-zA-Z0-9_]*)"
+                    match = re.match(pattern, full_word)
+                    if match:
+                        enum_name = match.group(1)
+                        variant_name = match.group(2)
+                        if enum_name in self.ast.get_enums():
+                            variants = self.ast.get_enum_variants(enum_name)
+                            if variant_name in variants:
+                                range = self.find_type_declaration(enum_name)
         if range:
             return range