diff --git a/tests/functional/grammar/test_grammar.py b/tests/functional/grammar/test_grammar.py index e95578e7e3..0ff8c23477 100644 --- a/tests/functional/grammar/test_grammar.py +++ b/tests/functional/grammar/test_grammar.py @@ -102,6 +102,6 @@ def has_no_docstrings(c): max_examples=500, suppress_health_check=[HealthCheck.too_slow, HealthCheck.filter_too_much] ) def test_grammar_bruteforce(code): - _, _, _, _, reformatted_code = pre_parse(code + "\n") - tree = parse_to_ast(reformatted_code) + pre_parse_result = pre_parse(code + "\n") + tree = parse_to_ast(pre_parse_result.reformatted_code) assert isinstance(tree, Module) diff --git a/tests/unit/ast/test_annotate_and_optimize_ast.py b/tests/unit/ast/test_annotate_and_optimize_ast.py index 67669c24c2..936c81e21a 100644 --- a/tests/unit/ast/test_annotate_and_optimize_ast.py +++ b/tests/unit/ast/test_annotate_and_optimize_ast.py @@ -28,20 +28,12 @@ def foo() -> int128: def get_contract_info(source_code): - ( - _, - loop_var_annotations, - native_hex_literal_locations, - class_types, - reformatted_code, - ) = pre_parse(source_code) - py_ast = python_ast.parse(reformatted_code) - - annotate_python_ast( - py_ast, reformatted_code, loop_var_annotations, native_hex_literal_locations, class_types - ) - - return py_ast, reformatted_code + pre_parse_result = pre_parse(source_code) + py_ast = python_ast.parse(pre_parse_result.reformatted_code) + + annotate_python_ast(py_ast, pre_parse_result) + + return py_ast, pre_parse_result.reformatted_code def test_it_annotates_ast_with_source_code(): diff --git a/vyper/ast/parse.py b/vyper/ast/parse.py index 0147492a11..1519217f52 100644 --- a/vyper/ast/parse.py +++ b/vyper/ast/parse.py @@ -6,7 +6,7 @@ import asttokens from vyper.ast import nodes as vy_ast -from vyper.ast.pre_parser import pre_parse +from vyper.ast.pre_parser import PreParseResult, pre_parse from vyper.compiler.settings import Settings from vyper.exceptions import CompilerPanic, ParserException, SyntaxException from vyper.typing import ModificationOffsets @@ -55,15 +55,9 @@ def parse_to_ast_with_settings( """ if "\x00" in vyper_source: raise ParserException("No null bytes (\\x00) allowed in the source code.") - ( - settings, - class_types, - for_loop_annotations, - native_hex_literal_locations, - python_source, - ) = pre_parse(vyper_source) + pre_parse_result = pre_parse(vyper_source) try: - py_ast = python_ast.parse(python_source) + py_ast = python_ast.parse(pre_parse_result.reformatted_code) except SyntaxError as e: # TODO: Ensure 1-to-1 match of source_code:reformatted_code SyntaxErrors raise SyntaxException(str(e), vyper_source, e.lineno, e.offset) from None @@ -78,23 +72,20 @@ def parse_to_ast_with_settings( annotate_python_ast( py_ast, - vyper_source, - class_types, - for_loop_annotations, - native_hex_literal_locations, + pre_parse_result, source_id=source_id, module_path=module_path, resolved_path=resolved_path, ) # postcondition: consumed all the for loop annotations - assert len(for_loop_annotations) == 0 + assert len(pre_parse_result.for_loop_annotations) == 0 # Convert to Vyper AST. module = vy_ast.get_node(py_ast) assert isinstance(module, vy_ast.Module) # mypy hint - return settings, module + return pre_parse_result.settings, module def ast_to_dict(ast_struct: Union[vy_ast.VyperNode, List]) -> Union[Dict, List]: @@ -124,10 +115,7 @@ def dict_to_ast(ast_struct: Union[Dict, List]) -> Union[vy_ast.VyperNode, List]: def annotate_python_ast( parsed_ast: python_ast.AST, - vyper_source: str, - modification_offsets: ModificationOffsets, - for_loop_annotations: dict, - native_hex_literal_locations: list, + pre_parse_result: PreParseResult, source_id: int = 0, module_path: Optional[str] = None, resolved_path: Optional[str] = None, @@ -139,30 +127,18 @@ def annotate_python_ast( ---------- parsed_ast : AST The AST to be annotated and optimized. - vyper_source: str - The original vyper source code - loop_var_annotations: dict - A mapping of line numbers of `For` nodes to the tokens of the type - annotation of the iterator extracted during pre-parsing. - modification_offsets : dict - A mapping of class names to their original class types. + pre_parse_result: PreParseResult + Outputs from pre-parsing. Returns ------- The annotated and optimized AST. """ - tokens = asttokens.ASTTokens(vyper_source) + tokens = asttokens.ASTTokens(pre_parse_result.reformatted_code) assert isinstance(parsed_ast, python_ast.Module) # help mypy tokens.mark_tokens(parsed_ast) visitor = AnnotatingVisitor( - vyper_source, - modification_offsets, - for_loop_annotations, - native_hex_literal_locations, - tokens, - source_id, - module_path=module_path, - resolved_path=resolved_path, + pre_parse_result, tokens, source_id, module_path=module_path, resolved_path=resolved_path ) visitor.visit(parsed_ast) @@ -176,10 +152,7 @@ class AnnotatingVisitor(python_ast.NodeTransformer): def __init__( self, - source_code: str, - modification_offsets: ModificationOffsets, - for_loop_annotations: dict, - native_hex_literal_locations: list, + pre_parse_result: PreParseResult, tokens: asttokens.ASTTokens, source_id: int, module_path: Optional[str] = None, @@ -189,10 +162,10 @@ def __init__( self._source_id = source_id self._module_path = module_path self._resolved_path = resolved_path - self._source_code = source_code - self._modification_offsets = modification_offsets - self._for_loop_annotations = for_loop_annotations - self._native_hex_literal_locations = native_hex_literal_locations + self._source_code = pre_parse_result.reformatted_code + self._modification_offsets = pre_parse_result.modification_offsets + self._for_loop_annotations = pre_parse_result.for_loop_annotations + self._native_hex_literal_locations = pre_parse_result.native_hex_literal_locations self.counter: int = 0 diff --git a/vyper/ast/pre_parser.py b/vyper/ast/pre_parser.py index 7e407f725c..d3f6e28ba2 100644 --- a/vyper/ast/pre_parser.py +++ b/vyper/ast/pre_parser.py @@ -12,7 +12,12 @@ # evm-version pragma from vyper.evm.opcodes import EVM_VERSIONS from vyper.exceptions import StructureException, SyntaxException, VersionException -from vyper.typing import ModificationOffsets, ParserPosition +from vyper.typing import ( + ForLoopAnnotations, + ModificationOffsets, + NativeHexLiteralLocations, + ParserPosition, +) def validate_version_pragma(version_str: str, full_source_code: str, start: ParserPosition) -> None: @@ -158,7 +163,34 @@ def consume(self, token, result): CUSTOM_EXPRESSION_TYPES = {"extcall": "ExtCall", "staticcall": "StaticCall"} -def pre_parse(code: str) -> tuple[Settings, ModificationOffsets, dict, list, str]: +class PreParseResult: + # Compilation settings based on the directives in the source code + settings: Settings + # A mapping of class names to their original class types. + modification_offsets: ModificationOffsets + # A mapping of line/column offsets of `For` nodes to the annotation of the for loop target + for_loop_annotations: ForLoopAnnotations + # A list of line/column offsets of native hex literals + native_hex_literal_locations: NativeHexLiteralLocations + # Reformatted python source string. + reformatted_code: str + + def __init__( + self, + settings, + modification_offsets, + for_loop_annotations, + native_hex_literal_locations, + reformatted_code, + ): + self.settings = settings + self.modification_offsets = modification_offsets + self.for_loop_annotations = for_loop_annotations + self.native_hex_literal_locations = native_hex_literal_locations + self.reformatted_code = reformatted_code + + +def pre_parse(code: str) -> PreParseResult: """ Re-formats a vyper source string into a python source string and performs some validation. More specifically, @@ -180,14 +212,8 @@ def pre_parse(code: str) -> tuple[Settings, ModificationOffsets, dict, list, str Returns ------- - Settings - Compilation settings based on the directives in the source code - ModificationOffsets - A mapping of class names to their original class types. - dict[tuple[int, int], list[TokenInfo]] - A mapping of line/column offsets of `For` nodes to the annotation of the for loop target - str - Reformatted python source string. + PreParseResult + Outputs for transforming the python AST to vyper AST """ result: list[TokenInfo] = [] modification_offsets: ModificationOffsets = {} @@ -311,7 +337,7 @@ def pre_parse(code: str) -> tuple[Settings, ModificationOffsets, dict, list, str for k, v in for_parser.annotations.items(): for_loop_annotations[k] = v.copy() - return ( + return PreParseResult( settings, modification_offsets, for_loop_annotations, diff --git a/vyper/typing.py b/vyper/typing.py index ad3964dff9..e01c149117 100644 --- a/vyper/typing.py +++ b/vyper/typing.py @@ -1,7 +1,10 @@ -from typing import Dict, Optional, Sequence, Tuple, Union +from tokenize import TokenInfo +from typing import Dict, Optional, List, Sequence, Tuple, Union # Parser +ForLoopAnnotations = Dict[Tuple[int, int], List[TokenInfo]] ModificationOffsets = Dict[Tuple[int, int], str] +NativeHexLiteralLocations = List[Tuple[int, int]] ParserPosition = Tuple[int, int] # Compiler