add raw ast output

vyperlang · Nov 15, 2023 · b0b39e4 · b0b39e4
1 parent cd5f6cc
commit b0b39e4
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 4 deletions.
diff --git a/vyper/cli/vyper_compile.py b/vyper/cli/vyper_compile.py
@@ -33,7 +33,8 @@
 devdoc             - Natspec developer documentation
 combined_json      - All of the above format options combined as single JSON output
 layout             - Storage layout of a Vyper contract
-ast                - AST in JSON format
+ast                - Annotated AST in JSON format
+raw_ast            - Unannotated AST in JSON format
 interface          - Vyper interface of a contract
 external_interface - External interface of a contract, used for outside contract calls
 opcodes            - List of opcodes as a string
@@ -239,7 +240,7 @@ def compile_files(
         output_formats = combined_json_outputs
         show_version = True
 
-    translate_map = {"abi_python": "abi", "json": "abi", "ast": "ast_dict", "ir_json": "ir_dict"}
+    translate_map = {"abi_python": "abi", "json": "abi", "ast": "ast_dict", "raw_ast": "raw_ast_dict", "ir_json": "ir_dict"}
     final_formats = [translate_map.get(i, i) for i in output_formats]
 
     if storage_layout_paths:

diff --git a/vyper/cli/vyper_json.py b/vyper/cli/vyper_json.py
@@ -323,6 +323,8 @@ def format_to_output_dict(compiler_data: dict) -> dict:
         output_dict["sources"][path] = {"id": data["source_id"]}
         if "ast_dict" in data:
             output_dict["sources"][path]["ast"] = data["ast_dict"]["ast"]
+        if "raw_ast_dict" in data:
+            output_dict["sources"][path]["raw_ast"] = data["raw_ast_dict"]["ast"]
 
         name = PurePath(path).stem
         output_dict["contracts"][path] = {name: {}}

diff --git a/vyper/compiler/__init__.py b/vyper/compiler/__init__.py
@@ -13,7 +13,8 @@
 
 OUTPUT_FORMATS = {
     # requires vyper_module
-    "ast_dict": output.build_ast_dict,
+    "raw_ast_dict": output.build_ast_dict,
+    "ast_dict": output.build_annotated_ast_dict,
     "layout": output.build_layout_output,
     # requires global_ctx
     "devdoc": output.build_devdoc,

diff --git a/vyper/compiler/output.py b/vyper/compiler/output.py
@@ -17,11 +17,19 @@
 def build_ast_dict(compiler_data: CompilerData) -> dict:
     ast_dict = {
         "contract_name": str(compiler_data.contract_path),
-        "ast": ast_to_dict(compiler_data.vyper_module_annotated),
+        "ast": ast_to_dict(compiler_data.vyper_module),
     }
     return ast_dict
 
 
+def build_annotated_ast_dict(compiler_data: CompilerData) -> dict:
+    annotated_ast_dict = {
+        "contract_name": str(compiler_data.contract_path),
+        "ast": ast_to_dict(compiler_data.vyper_module_annotated),
+    }
+    return annotated_ast_dict
+
+
 def build_devdoc(compiler_data: CompilerData) -> dict:
     userdoc, devdoc = parse_natspec(compiler_data.vyper_module_folded)
     return devdoc