feat: require type annotations for loop variables

vyper/ast/parse.py

@@ -12,9 +12,9 @@
 from vyper.typing import ModificationOffsets
 
 
-def parse_to_ast(*args: Any, **kwargs: Any) -> vy_ast.Module:
-    _settings, ast = parse_to_ast_with_settings(*args, **kwargs)
-    return ast
+def parse_to_ast(*args: Any, **kwargs: Any) -> tuple[vy_ast.Module, dict[int, dict[str, Any]]]:
+    _settings, ast, loop_var_annotations = parse_to_ast_with_settings(*args, **kwargs)
+    return ast, loop_var_annotations
 
 
 def parse_to_ast_with_settings(
@@ -23,7 +23,7 @@ def parse_to_ast_with_settings(
     module_path: Optional[str] = None,
     resolved_path: Optional[str] = None,
     add_fn_node: Optional[str] = None,
-) -> tuple[Settings, vy_ast.Module]:
+) -> tuple[Settings, vy_ast.Module, dict[int, dict[str, Any]]]:
     """
     Parses a Vyper source string and generates basic Vyper AST nodes.
 
@@ -54,7 +54,7 @@ def parse_to_ast_with_settings(
     """
     if "\x00" in source_code:
         raise ParserException("No null bytes (\\x00) allowed in the source code.")
-    settings, class_types, reformatted_code = pre_parse(source_code)
+    settings, class_types, loop_var_annotations, reformatted_code = pre_parse(source_code)
     try:
         py_ast = python_ast.parse(reformatted_code)
     except SyntaxError as e:
@@ -73,6 +73,7 @@ def parse_to_ast_with_settings(
         py_ast,
         source_code,
         class_types,
+        loop_var_annotations,
         source_id,
         module_path=module_path,
         resolved_path=resolved_path,
@@ -82,7 +83,7 @@ def parse_to_ast_with_settings(
     module = vy_ast.get_node(py_ast)
     assert isinstance(module, vy_ast.Module)  # mypy hint
 
-    return settings, module
+    return settings, module, loop_var_annotations
 
 
 def ast_to_dict(ast_struct: Union[vy_ast.VyperNode, List]) -> Union[Dict, List]:
@@ -356,6 +357,7 @@ def annotate_python_ast(
     parsed_ast: python_ast.AST,
     source_code: str,
     modification_offsets: Optional[ModificationOffsets] = None,
+    loop_var_annotations: Optional[dict[int, python_ast.AST]] = None,
     source_id: int = 0,
     module_path: Optional[str] = None,
     resolved_path: Optional[str] = None,
@@ -387,5 +389,9 @@ def annotate_python_ast(
         resolved_path=resolved_path,
     )
     visitor.visit(parsed_ast)
+    for k, v in loop_var_annotations.items():
+        tokens = asttokens.ASTTokens(v["source_code"], tree=cast(Optional[python_ast.Module], v["parsed_ast"]))
+        visitor = AnnotatingVisitor(v["source_code"], {}, tokens, source_id, contract_name)
+        visitor.visit(v["parsed_ast"])
 
     return parsed_ast

vyper/ast/pre_parser.py

@@ -79,11 +79,16 @@ def pre_parse(code: str) -> tuple[Settings, ModificationOffsets, str]:
     result = []
     modification_offsets: ModificationOffsets = {}
     settings = Settings()
+    loop_var_annotations = {}
 
     try:
         code_bytes = code.encode("utf-8")
         token_list = list(tokenize(io.BytesIO(code_bytes).readline))
 
+        is_for_loop = False
+        after_loop_var = False
+        loop_var_annotation = []
+
         for i in range(len(token_list)):
             token = token_list[i]
             toks = [token]
@@ -146,8 +151,44 @@ def pre_parse(code: str) -> tuple[Settings, ModificationOffsets, str]:
 
             if (typ, string) == (OP, ";"):
                 raise SyntaxException("Semi-colon statements not allowed", code, start[0], start[1])
+
+            if typ == NAME and string == "for":
+                is_for_loop = True
+                #print("for loop!")
+                #print(token)
+
+            if is_for_loop:
+                if typ == NAME and string == "in":
+                    loop_var_annotations[start[0]] = loop_var_annotation
+
+                    is_for_loop = False
+                    after_loop_var = False
+                    loop_var_annotation = []
+
+                elif (typ, string) == (OP, ":"):
+                    after_loop_var = True
+                    continue
+
+                elif after_loop_var and not (typ == NAME and string == "for"):
+                    #print("adding to loop var: ", toks)
+                    loop_var_annotation.extend(toks)
+                    continue
+
             result.extend(toks)
     except TokenError as e:
         raise SyntaxException(e.args[0], code, e.args[1][0], e.args[1][1]) from e
 
-    return settings, modification_offsets, untokenize(result).decode("utf-8")
+    for k, v in loop_var_annotations.items():
+
+
+        updated_v = untokenize(v)
+        #print("untokenized v: ", updated_v)
+        updated_v = updated_v.replace("\\", "")
+        updated_v = updated_v.replace("\n", "")
+        import textwrap
+        #print("updated v: ", textwrap.dedent(updated_v))
+        loop_var_annotations[k] = {"source_code": textwrap.dedent(updated_v)}
+
+    #print("untokenized result: ", type(untokenize(result)))
+    #print("untokenized result decoded: ", untokenize(result).decode("utf-8"))
+    return settings, modification_offsets, loop_var_annotations, untokenize(result).decode("utf-8")

vyper/compiler/phases.py

@@ -2,7 +2,7 @@
 import warnings
 from functools import cached_property
 from pathlib import Path, PurePath
-from typing import Optional
+from typing import Any, Optional
 
 from vyper import ast as vy_ast
 from vyper.codegen import module
@@ -129,7 +129,7 @@ def contract_path(self):
 
     @cached_property
     def _generate_ast(self):
-        settings, ast = vy_ast.parse_to_ast_with_settings(
+        settings, ast, loop_var_annotations = vy_ast.parse_to_ast_with_settings(
             self.source_code,
             self.source_id,
             module_path=str(self.contract_path),
@@ -145,16 +145,17 @@ def _generate_ast(self):
 
         # note self.settings.compiler_version is erased here as it is
         # not used after pre-parsing
-        return ast
+        return ast, loop_var_annotations
 
     @cached_property
     def vyper_module(self):
         return self._generate_ast
 
     @cached_property
     def _annotated_module(self):
+        ast, loop_var_annotations = self.vyper_module
         return generate_annotated_ast(
-            self.vyper_module, self.input_bundle, self.storage_layout_override
+            ast, loop_var_annotations, self.input_bundle, self.storage_layout_override
         )
 
     @property
@@ -244,6 +245,7 @@ def blueprint_bytecode(self) -> bytes:
 
 def generate_annotated_ast(
     vyper_module: vy_ast.Module,
+    loop_var_annotations: dict[int, dict[str, Any]],
     input_bundle: InputBundle,
     storage_layout_overrides: StorageLayout = None,
 ) -> tuple[vy_ast.Module, StorageLayout]:
@@ -265,7 +267,7 @@ def generate_annotated_ast(
     vyper_module = copy.deepcopy(vyper_module)
     with input_bundle.search_path(Path(vyper_module.resolved_path).parent):
         # note: validate_semantics does type inference on the AST
-        validate_semantics(vyper_module, input_bundle)
+        validate_semantics(vyper_module, loop_var_annotations, input_bundle)
 
     symbol_tables = set_data_positions(vyper_module, storage_layout_overrides)
 

vyper/semantics/analysis/local.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Any, Optional
 
 from vyper import ast as vy_ast
 from vyper.ast.metadata import NodeMetadata
@@ -53,15 +53,15 @@
 from vyper.semantics.types.utils import type_from_annotation
 
 
-def validate_functions(vy_module: vy_ast.Module) -> None:
+def validate_functions(vy_module: vy_ast.Module, loop_var_annotations: dict[int, dict[str, Any]]) -> None:
     """Analyzes a vyper ast and validates the function bodies"""
 
     err_list = ExceptionList()
     namespace = get_namespace()
     for node in vy_module.get_children(vy_ast.FunctionDef):
         with namespace.enter_scope():
             try:
-                analyzer = FunctionNodeVisitor(vy_module, node, namespace)
+                analyzer = FunctionNodeVisitor(vy_module, loop_var_annotations, node, namespace)
                 analyzer.analyze()
             except VyperException as e:
                 err_list.append(e)
@@ -180,9 +180,10 @@
     scope_name = "function"
 
     def __init__(
-        self, vyper_module: vy_ast.Module, fn_node: vy_ast.FunctionDef, namespace: dict
+        self, vyper_module: vy_ast.Module, loop_var_annotations: dict[int, dict[str, Any]], fn_node: vy_ast.FunctionDef, namespace: dict
     ) -> None:
         self.vyper_module = vyper_module
+        self.loop_var_annotations = loop_var_annotations
         self.fn_node = fn_node
         self.namespace = namespace
         self.func = fn_node._metadata["func_type"]
@@ -350,6 +351,11 @@
         if isinstance(node.iter, vy_ast.Subscript):
             raise StructureException("Cannot iterate over a nested list", node.iter)
 
+        print("visit_For: ", node.lineno)
+        iter_type = type_from_annotation(self.loop_var_annotations[node.lineno]["vy_ast"].body[0].value)
+        print("iter type: ", iter_type)
+        node.target._metadata["type"] = iter_type   
+
         if isinstance(node.iter, vy_ast.Call):
             # iteration via range()
             if node.iter.get("func.id") != "range":
@@ -418,62 +424,62 @@
         if not isinstance(node.target, vy_ast.Name):
             raise StructureException("Invalid syntax for loop iterator", node.target)
 
-        for_loop_exceptions = []
-        iter_name = node.target.id
-        for possible_target_type in type_list:
-            # type check the for loop body using each possible type for iterator value
-
-            with self.namespace.enter_scope():
-                self.namespace[iter_name] = VarInfo(
-                    possible_target_type, modifiability=Modifiability.RUNTIME_CONSTANT
-                )
-
-                try:
-                    with NodeMetadata.enter_typechecker_speculation():
-                        for stmt in node.body:
-                            self.visit(stmt)
-
-                        self.expr_visitor.visit(node.target, possible_target_type)
-
-                        if isinstance(node.iter, (vy_ast.Name, vy_ast.Attribute)):
-                            iter_type = get_exact_type_from_node(node.iter)
-                            # note CMC 2023-10-23: slightly redundant with how type_list is computed
-                            validate_expected_type(node.target, iter_type.value_type)
-                            self.expr_visitor.visit(node.iter, iter_type)
-                        if isinstance(node.iter, vy_ast.List):
-                            len_ = len(node.iter.elements)
-                            self.expr_visitor.visit(node.iter, SArrayT(possible_target_type, len_))
-                        if isinstance(node.iter, vy_ast.Call) and node.iter.func.id == "range":
-                            for a in node.iter.args:
-                                self.expr_visitor.visit(a, possible_target_type)
-                            for a in node.iter.keywords:
-                                if a.arg == "bound":
-                                    self.expr_visitor.visit(a.value, possible_target_type)
-
-                except (TypeMismatch, InvalidOperation) as exc:
-                    for_loop_exceptions.append(exc)
-                else:
-                    # success -- do not enter error handling section
-                    return
-
-        # failed to find a good type. bail out
-        if len(set(str(i) for i in for_loop_exceptions)) == 1:
-            # if every attempt at type checking raised the same exception
-            raise for_loop_exceptions[0]
-
-        # return an aggregate TypeMismatch that shows all possible exceptions
-        # depending on which type is used
-        types_str = [str(i) for i in type_list]
-        given_str = f"{', '.join(types_str[:1])} or {types_str[-1]}"
-        raise TypeMismatch(
-            f"Iterator value '{iter_name}' may be cast as {given_str}, "
-            "but type checking fails with all possible types:",
-            node,
-            *(
-                (f"Casting '{iter_name}' as {typ}: {exc.message}", exc.annotations[0])
-                for typ, exc in zip(type_list, for_loop_exceptions)
-            ),
-        )
+        # for_loop_exceptions = []
+        # iter_name = node.target.id
+        # for possible_target_type in type_list:
+        #     # type check the for loop body using each possible type for iterator value
+
+        #     with self.namespace.enter_scope():
+        #         self.namespace[iter_name] = VarInfo(
+        #             possible_target_type, modifiability=Modifiability.RUNTIME_CONSTANT
+        #         )
+
+        #         try:
+        #             with NodeMetadata.enter_typechecker_speculation():
+        #                 for stmt in node.body:
+        #                     self.visit(stmt)
+
+        #                 self.expr_visitor.visit(node.target, possible_target_type)
+
+        #                 if isinstance(node.iter, (vy_ast.Name, vy_ast.Attribute)):
+        #                     iter_type = get_exact_type_from_node(node.iter)
+        #                     # note CMC 2023-10-23: slightly redundant with how type_list is computed
+        #                     validate_expected_type(node.target, iter_type.value_type)
+        #                     self.expr_visitor.visit(node.iter, iter_type)
+        #                 if isinstance(node.iter, vy_ast.List):
+        #                     len_ = len(node.iter.elements)
+        #                     self.expr_visitor.visit(node.iter, SArrayT(possible_target_type, len_))
+        #                 if isinstance(node.iter, vy_ast.Call) and node.iter.func.id == "range":
+        #                     for a in node.iter.args:
+        #                         self.expr_visitor.visit(a, possible_target_type)
+        #                     for a in node.iter.keywords:
+        #                         if a.arg == "bound":
+        #                             self.expr_visitor.visit(a.value, possible_target_type)
+
+        #         except (TypeMismatch, InvalidOperation) as exc:
+        #             for_loop_exceptions.append(exc)
+        #         else:
+        #             # success -- do not enter error handling section
+        #             return
+
+        # # failed to find a good type. bail out
+        # if len(set(str(i) for i in for_loop_exceptions)) == 1:
+        #     # if every attempt at type checking raised the same exception
+        #     raise for_loop_exceptions[0]
+
+        # # return an aggregate TypeMismatch that shows all possible exceptions
+        # # depending on which type is used
+        # types_str = [str(i) for i in type_list]
+        # given_str = f"{', '.join(types_str[:1])} or {types_str[-1]}"
+        # raise TypeMismatch(
+        #     f"Iterator value '{iter_name}' may be cast as {given_str}, "
+        #     "but type checking fails with all possible types:",
+        #     node,
+        #     *(
+        #         (f"Casting '{iter_name}' as {typ}: {exc.message}", exc.annotations[0])
+        #         for typ, exc in zip(type_list, for_loop_exceptions)
+        #     ),
+        # )
 
     def visit_If(self, node):
         validate_expected_type(node.test, BoolT())