asdl_rs.py to multiple file output

Author: youknowone

ast/asdl_rs.py

@@ -12,7 +12,7 @@
 import asdl
 
 TABSIZE = 4
-AUTOGEN_MESSAGE = "// File automatically generated by {}.\n"
+AUTOGEN_MESSAGE = "// File automatically generated by {}.\n\n"
 
 builtin_type_mapping = {
     "identifier": "Ident",
@@ -535,13 +535,9 @@ def gen_construction(self, header, fields, footer, depth):
 class FoldModuleVisitor(EmitVisitor):
     def visitModule(self, mod):
         depth = 0
-        self.emit('#[cfg(feature = "fold")]', depth)
-        self.emit("pub mod fold {", depth)
-        self.emit("use super::*;", depth + 1)
-        self.emit("use crate::fold_helpers::Foldable;", depth + 1)
-        FoldTraitDefVisitor(self.file, self.typeinfo).visit(mod, depth + 1)
-        FoldImplVisitor(self.file, self.typeinfo).visit(mod, depth + 1)
-        self.emit("}", depth)
+        self.emit("use crate::fold_helpers::Foldable;", depth)
+        FoldTraitDefVisitor(self.file, self.typeinfo).visit(mod, depth)
+        FoldImplVisitor(self.file, self.typeinfo).visit(mod, depth)
 
 
 class VisitorTraitDefVisitor(StructVisitor):
@@ -577,7 +573,9 @@ def emit_visitor(self, nodename, depth, has_node=True):
         else:
             node_type = "()"
             node_value = "()"
-        self.emit(f"fn visit_{typeinfo.sum_name}(&mut self, node: {node_type}) {{", depth)
+        self.emit(
+            f"fn visit_{typeinfo.sum_name}(&mut self, node: {node_type}) {{", depth
+        )
         self.emit(f"self.generic_visit_{typeinfo.sum_name}({node_value})", depth + 1)
         self.emit("}", depth)
 
@@ -587,7 +585,10 @@ def emit_generic_visitor_signature(self, nodename, depth, has_node=True):
             node_type = typeinfo.rust_sum_name
         else:
             node_type = "()"
-        self.emit(f"fn generic_visit_{typeinfo.sum_name}(&mut self, node: {node_type}) {{", depth)
+        self.emit(
+            f"fn generic_visit_{typeinfo.sum_name}(&mut self, node: {node_type}) {{",
+            depth,
+        )
 
     def emit_empty_generic_visitor(self, nodename, depth):
         self.emit_generic_visitor_signature(nodename, depth)
@@ -665,13 +666,8 @@ def visitProduct(self, product, name, depth):
 class VisitorModuleVisitor(EmitVisitor):
     def visitModule(self, mod):
         depth = 0
-        self.emit('#[cfg(feature = "visitor")]', depth)
         self.emit("#[allow(unused_variables, non_snake_case)]", depth)
-        self.emit("pub mod visitor {", depth)
-        self.emit("use super::*;", depth + 1)
-        VisitorTraitDefVisitor(self.file, self.typeinfo).visit(mod, depth + 1)
-        self.emit("}", depth)
-        self.emit("", depth)
+        VisitorTraitDefVisitor(self.file, self.typeinfo).visit(mod, depth)
 
 
 class ClassDefVisitor(EmitVisitor):
@@ -930,27 +926,19 @@ def visit(self, object):
             v.emit("", 0)
 
 
-def write_generic_def(mod, typeinfo, f):
-    f.write(
-        textwrap.dedent(
-            """
-        pub use crate::{Attributed, constant::*};
+def write_ast_def(mod, typeinfo, f):
+    StructVisitor(f, typeinfo).visit(mod)
 
-        type Ident = String;
-        \n
-        """
-        )
-    )
 
-    c = ChainOfVisitors(
-        StructVisitor(f, typeinfo),
-        FoldModuleVisitor(f, typeinfo),
-        VisitorModuleVisitor(f, typeinfo),
-    )
-    c.visit(mod)
+def write_fold_def(mod, typeinfo, f):
+    FoldModuleVisitor(f, typeinfo).visit(mod)
+
+
+def write_visitor_def(mod, typeinfo, f):
+    VisitorModuleVisitor(f, typeinfo).visit(mod)
 
 
-def write_located_def(typeinfo, f):
+def write_located_def(mod, typeinfo, f):
     f.write(
         textwrap.dedent(
             """
@@ -1000,8 +988,7 @@ def write_ast_mod(mod, typeinfo, f):
 
 def main(
     input_filename,
-    generic_filename,
-    located_filename,
+    ast_dir,
     module_filename,
     dump_module=False,
 ):
@@ -1016,34 +1003,34 @@ def main(
     typeinfo = {}
     FindUserdataTypesVisitor(typeinfo).visit(mod)
 
-    with generic_filename.open("w") as generic_file, located_filename.open(
-        "w"
-    ) as located_file:
-        generic_file.write(auto_gen_msg)
-        write_generic_def(mod, typeinfo, generic_file)
-        located_file.write(auto_gen_msg)
-        write_located_def(typeinfo, located_file)
+    for filename, write in [
+        ("generic", write_ast_def),
+        ("fold", write_fold_def),
+        ("located", write_located_def),
+        ("visitor", write_visitor_def),
+    ]:
+        with (ast_dir / f"{filename}.rs").open("w") as f:
+            f.write(auto_gen_msg)
+            write(mod, typeinfo, f)
 
     with module_filename.open("w") as module_file:
         module_file.write(auto_gen_msg)
         write_ast_mod(mod, typeinfo, module_file)
 
-    print(f"{generic_filename}, {located_filename}, {module_filename} regenerated.")
+    print(f"{ast_dir}, {module_filename} regenerated.")
 
 
 if __name__ == "__main__":
     parser = ArgumentParser()
     parser.add_argument("input_file", type=Path)
-    parser.add_argument("-G", "--generic-file", type=Path, required=True)
-    parser.add_argument("-L", "--located-file", type=Path, required=True)
+    parser.add_argument("-A", "--ast-dir", type=Path, required=True)
     parser.add_argument("-M", "--module-file", type=Path, required=True)
     parser.add_argument("-d", "--dump-module", action="store_true")
 
     args = parser.parse_args()
     main(
         args.input_file,
-        args.generic_file,
-        args.located_file,
+        args.ast_dir,
         args.module_file,
         args.dump_module,
     )