Reject untyped NA variable declarations

Author: Poryaei

src/pinescript_validator/ast_validator.py

@@ -198,6 +198,8 @@ def validate_statement(self, statement: AST.Statement, scope: Scope, conditional
                 self.validate_type_annotation(statement.type_annotation, scope, statement.line, statement.column)
             if statement.init is not None:
                 self.validate_expression(statement.init, scope, conditional_context)
+                if statement.type_annotation is None and self.is_na_literal(statement.init):
+                    self.untyped_na_declaration_error(statement.name_line, statement.name_column, statement.name)
                 if self.is_bool_type_annotation(statement.type_annotation) and self.is_na_literal(statement.init):
                     self.bool_na_error(statement.init.line, statement.init.column)
             return
@@ -1281,6 +1283,18 @@ def bool_na_error(self, line: int, column: int) -> None:
             )
         )
 
+    def untyped_na_declaration_error(self, line: int, column: int, name: str) -> None:
+        self.errors.append(
+            Diagnostic(
+                line=line,
+                column=column,
+                length=len(name),
+                message="Value with NA type cannot be assigned to a variable that was defined without type keyword",
+                severity=Severity.ERROR,
+                source="ast",
+            )
+        )
+
     def lookup_accessible_symbol(
         self,
         scope: Scope,

tests/test_validator.py

@@ -50,6 +50,24 @@ def test_unused_variable(self) -> None:
         diagnostics = self.validator.validate_text('indicator("Test")\nvalue = close')
         self.assertTrue(any("never used" in diagnostic.message for diagnostic in diagnostics))
 
+    def test_untyped_variable_declaration_cannot_use_na(self) -> None:
+        diagnostics = self.validator.validate_text("x = na")
+        self.assertTrue(
+            any(
+                "Value with NA type cannot be assigned to a variable that was defined without type keyword" in diagnostic.message
+                for diagnostic in diagnostics
+            )
+        )
+
+    def test_typed_variable_declaration_can_use_na(self) -> None:
+        diagnostics = self.validator.validate_text("float x = na")
+        self.assertFalse(
+            any(
+                "Value with NA type cannot be assigned to a variable that was defined without type keyword" in diagnostic.message
+                for diagnostic in diagnostics
+            )
+        )
+
     def test_invalid_named_argument(self) -> None:
         diagnostics = self.validator.validate_text('plot(close, invalid_param=true)')
         self.assertTrue(any("Invalid parameter 'invalid_param'" in diagnostic.message for diagnostic in diagnostics))