Add gen mod for more convenient to use

while use customize ,you can use gen_mod to generate mod.rs.

Signed-off-by: jokemanfire <[email protected]>

.gitignore

@@ -6,5 +6,4 @@ Cargo.lock
 .idea
 *.o
 example/protocols/**/*.rs
-!example/protocols/**/mod.rs
 src/ttrpc.rs

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::{
@@ -723,6 +727,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);

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,
 }

example/build.rs

@@ -4,12 +4,15 @@
 //
 
 use std::{
-    fs::File,
+    fs::{self, File},
     io::{Read, Write},
 };
 use ttrpc_codegen::{Codegen, Customize, ProtobufCustomize};
 
 fn main() {
+    fs::create_dir_all("protocols/sync").unwrap();
+    fs::create_dir_all("protocols/asynchronous").unwrap();
+
     let mut protos = vec![
         "protocols/protos/github.com/gogo/protobuf/gogoproto/gogo.proto",
         "protocols/protos/github.com/kata-containers/agent/pkg/types/types.proto",
@@ -18,15 +21,15 @@ fn main() {
         "protocols/protos/google/protobuf/empty.proto",
         "protocols/protos/oci.proto",
     ];
-
-    let protobuf_customized = ProtobufCustomize::default().gen_mod_rs(false);
+    let protobuf_customized = ProtobufCustomize::default().gen_mod_rs(true);
 
     Codegen::new()
         .out_dir("protocols/sync")
         .inputs(&protos)
         .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 +45,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()
         })