Fork/charles cooper/feat/export ast

tests/unit/ast/test_ast_dict.py

@@ -1780,3 +1780,48 @@ def qux2():
             },
         }
     ]
+
+def test_annotated_ast_export_recursion(make_input_bundle):
+    sources = {
+        "main.vy": """
+import lib1
+
+@external
+def foo():
+    lib1.foo()
+    """,
+        "lib1.vy": """
+import lib2
+
+def foo():
+    lib2.foo()
+    """,
+        "lib2.vy": """
+def foo():
+    pass
+    """
+    }
+
+    input_bundle = make_input_bundle(sources)
+
+    def compile_and_get_ast(file_name):
+        file = input_bundle.load_file(file_name)
+        output = compiler.compile_from_file_input(
+            file, input_bundle=input_bundle, output_formats=["annotated_ast_dict"]
+        )
+        return output["annotated_ast_dict"]
+
+    lib1_ast = compile_and_get_ast("lib1.vy")["ast"]
+    lib2_ast = compile_and_get_ast("lib2.vy")["ast"]
+    main_out = compile_and_get_ast("main.vy")
+
+    lib1_import_ast = main_out["imports"][1]
+    lib2_import_ast = main_out["imports"][0]
+
+    # path is once virtual, once libX.vy
+    # type contains name which is based on path
+    kws = [s for s in lib1_import_ast.keys() if s not in {"path", "type"}]
+
+    for kw in kws:
+        assert lib1_ast[kw] == lib1_import_ast[kw]
+        assert lib2_ast[kw] == lib2_import_ast[kw]

vyper/compiler/output.py

@@ -30,6 +30,7 @@ def build_ast_dict(compiler_data: CompilerData) -> dict:
 
 def build_annotated_ast_dict(compiler_data: CompilerData) -> dict:
     module_t = compiler_data.annotated_vyper_module._metadata["type"]
+    # get all reachable imports including recursion
     imported_module_infos = module_t.reachable_imports
     unique_modules: dict[str, vy_ast.Module] = {}
     for info in imported_module_infos:
@@ -46,7 +47,7 @@ def build_annotated_ast_dict(compiler_data: CompilerData) -> dict:
         # come from multiple InputBundles (particularly builtin interfaces),
         # so source_id is not guaranteed to be unique.
         if ast.resolved_path in unique_modules:
-            # sanity check -- object equality
+            # sanity check -- objects must be identical
             assert unique_modules[ast.resolved_path] is ast
         unique_modules[ast.resolved_path] = ast