replace EnumDef with FlagDef

vyperlang · Aug 8, 2024 · de9763e · de9763e
1 parent 1ee9b66
commit de9763e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/vyper_lsp/ast.py b/vyper_lsp/ast.py
@@ -13,7 +13,7 @@ class AST:
     ast_data_folded = None
     ast_data_unfolded = None
 
-    custom_type_node_types = (nodes.StructDef, nodes.EnumDef)
+    custom_type_node_types = (nodes.StructDef, nodes.FlagDef)
 
     @classmethod
     def from_node(cls, node: VyperNode):
@@ -67,7 +67,7 @@ def get_top_level_nodes(self, *args, **kwargs):
         return self.best_ast.get_children(*args, **kwargs)
 
     def get_enums(self) -> List[str]:
-        return [node.name for node in self.get_descendants(nodes.EnumDef)]
+        return [node.name for node in self.get_descendants(nodes.FlagDef)]
 
     def get_structs(self) -> List[str]:
         return [node.name for node in self.get_descendants(nodes.StructDef)]
@@ -153,7 +153,7 @@ def get_attributes_for_symbol(self, symbol: str):
 
         if isinstance(node, nodes.StructDef):
             return self.get_struct_fields(symbol)
-        elif isinstance(node, nodes.EnumDef):
+        elif isinstance(node, nodes.FlagDef):
             return self.get_enum_variants(symbol)
 
         return []
@@ -186,7 +186,7 @@ def find_type_declaration_node_for_name(self, symbol: str):
         for node in self.get_descendants(searchable_types):
             if node.name == symbol:
                 return node
-            if isinstance(node, nodes.EnumDef):
+            if isinstance(node, nodes.FlagDef):
                 for variant in node.get_children(nodes.Expr):
                     if variant.value.id == symbol:
                         return variant

diff --git a/vyper_lsp/navigation.py b/vyper_lsp/navigation.py
@@ -4,7 +4,7 @@
 from typing import List, Optional
 
 from pygls.workspace import Document
-from vyper.ast import EnumDef, FunctionDef, VyperNode
+from vyper.ast import FlagDef, FunctionDef, VyperNode
 from vyper_lsp.ast import AST
 from vyper_lsp.utils import (
     get_expression_at_cursor,
@@ -98,7 +98,7 @@ def finalize(refs):
         if self._is_state_var_decl(og_line, word):
             return finalize(self.ast.find_nodes_referencing_state_variable(word))
 
-        if isinstance(top_level_node, EnumDef):
+        if isinstance(top_level_node, FlagDef):
             return finalize(
                 self.ast.find_nodes_referencing_enum_variant(top_level_node.name, word)
             )