diff --git a/.gitignore b/.gitignore index 14dabae..ff70e92 100644 --- a/.gitignore +++ b/.gitignore @@ -6,5 +6,4 @@ Cargo.lock .idea *.o example/protocols/**/*.rs -!example/protocols/**/mod.rs src/ttrpc.rs diff --git a/compiler/src/codegen.rs b/compiler/src/codegen.rs index 382f5ae..5c87ed8 100644 --- a/compiler/src/codegen.rs +++ b/compiler/src/codegen.rs @@ -36,7 +36,11 @@ #![allow(dead_code)] -use std::collections::HashMap; +use std::{ + collections::{HashMap, HashSet}, + fs, + io::BufRead, +}; use crate::Customize; use protobuf::{ @@ -724,6 +728,31 @@ pub fn gen_and_write( ) -> io::Result<()> { let results = gen(file_descriptors, files_to_generate, customize); + if customize.gen_mod { + let file_path = out_dir.join("mod.rs"); + let mut set = HashSet::new(); + //if mod file exists + if let Ok(file) = File::open(&file_path) { + let reader = io::BufReader::new(file); + reader.lines().for_each(|line| { + let _ = line.map(|r| set.insert(r)); + }); + } + let mut file_write = fs::OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(&file_path)?; + for r in &results { + let prefix_name: Vec<&str> = r.name.split('.').collect(); + set.insert(format!("pub mod {};", prefix_name[0])); + } + for item in &set { + writeln!(file_write, "{}", item)?; + } + file_write.flush()?; + } + for r in &results { let mut file_path = out_dir.to_owned(); file_path.push(&r.name); diff --git a/compiler/src/lib.rs b/compiler/src/lib.rs index 25a84e0..444fc61 100644 --- a/compiler/src/lib.rs +++ b/compiler/src/lib.rs @@ -36,4 +36,6 @@ pub struct Customize { pub async_client: bool, /// Indicates whether to generate async code for server. pub async_server: bool, + /// Gen mod rs in mod.rs + pub gen_mod: bool, } diff --git a/example/build.rs b/example/build.rs index a90274b..5026b4b 100644 --- a/example/build.rs +++ b/example/build.rs @@ -18,8 +18,9 @@ fn main() { "protocols/protos/google/protobuf/empty.proto", "protocols/protos/oci.proto", ]; - - let protobuf_customized = ProtobufCustomize::default().gen_mod_rs(false); + ///This could be add while the new ttrpc compiler support it. + ///If need use now change ttrpc_codegen's cargo.toml ttrpc-compiler = "0.6.1" to ttrpc-compiler = "../compiler" + let protobuf_customized = ProtobufCustomize::default().gen_mod_rs(true); Codegen::new() .out_dir("protocols/sync") @@ -27,6 +28,7 @@ fn main() { .include("protocols/protos") .rust_protobuf() .customize(Customize { + gen_mod:true,//This could be add while the new ttrpc compiler support it. ..Default::default() }) .rust_protobuf_customize(protobuf_customized.clone()) @@ -42,6 +44,7 @@ fn main() { .include("protocols/protos") .rust_protobuf() .customize(Customize { + gen_mod:true, //This could be add while the new ttrpc compiler support it. async_all: true, ..Default::default() })