move AnnAssign validation to vyper/ast/nodes.py

charles-cooper · charles-cooper · commit 73a322895e2c · 2024-02-24T15:30:15.000Z
diff --git a/vyper/ast/nodes.py b/vyper/ast/nodes.py
@@ -1261,7 +1261,16 @@ def __init__(self, *args, **kwargs):
 
 
 class AnnAssign(VyperNode):
-    __slots__ = ("target", "annotation", "value", "simple")
+    __slots__ = ("target", "annotation", "value")
+
+    def validate(self):
+        if not isinstance(self.target, Name):
+            raise VariableDeclarationException("Invalid variable declaration", self.target)
+
+        if self.value is None:
+            raise VariableDeclarationException(
+                "Local variables must be declared with an initial value", self
+            )
 
 
 class VariableDecl(VyperNode):
diff --git a/vyper/ast/nodes.pyi b/vyper/ast/nodes.pyi
@@ -208,7 +208,7 @@ class Assign(VyperNode): ...
 
 class AnnAssign(VyperNode):
     target: Name = ...
-    value: VyperNode = ...
+    value: ExprNode = ...
     annotation: VyperNode = ...
 
 class VariableDecl(VyperNode):
diff --git a/vyper/semantics/analysis/local.py b/vyper/semantics/analysis/local.py
@@ -17,7 +17,6 @@
     StructureException,
     TypeCheckFailure,
     TypeMismatch,
-    VariableDeclarationException,
     VyperException,
 )
 from vyper.semantics.analysis.base import (
@@ -326,14 +325,11 @@ def visit(self, node):
         super().visit(node)
 
     def visit_AnnAssign(self, node):
-        name = node.get("target.id")
-        if name is None:
-            raise VariableDeclarationException("Invalid assignment", node)
+        name = node.target.id
 
-        if not node.value:
-            raise VariableDeclarationException(
-                "Memory variables must be declared with an initial value", node
-            )
+        # sanity check postconditions of AnnAssign.validate()
+        assert isinstance(node.target, vy_ast.Name)
+        assert node.value is not None
 
         typ = type_from_annotation(node.annotation, DataLocation.MEMORY)
 
diff --git a/vyper/semantics/analysis/module.py b/vyper/semantics/analysis/module.py
@@ -20,7 +20,6 @@
     StateAccessViolation,
     StructureException,
     UndeclaredDefinition,
-    VariableDeclarationException,
     VyperException,
     tag_exceptions,
 )

Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,6 @@`
`20`	`20`	`StateAccessViolation,`
`21`	`21`	`StructureException,`
`22`	`22`	`UndeclaredDefinition,`
`23`		`- VariableDeclarationException,`
`24`	`23`	`VyperException,`
`25`	`24`	`tag_exceptions,`
`26`	`25`	`)`