build.rs

use std::fs;
use std::fs::File;
use std::io::prelude::*;
use ttrpc_codegen::Codegen;
use ttrpc_codegen::{Customize, ProtobufCustomize};

fn main() {
    let mut protos = fs::read_dir("protos")
        .expect("Could not read 'protos' directory")
        .map(|res| {
            res.map(|e| e.path())
                .expect("Could not get path of proto files")
        })
        .collect::<Vec<_>>();

    // Sort entries
    protos.sort();

    // Create the lib.rs file and add all the auto-generated mods
    let mut libf = File::create("src/lib.rs").expect("Could not open crate's lib file");
    libf.write(b"// This is auto-generated by build.rs\n\n")
        .expect("Could not write in crate's lib.rs");

    for proto in protos.iter() {
        // Re-run compilation every time the proto file changes
        println!("cargo:rerun-if-changed={:?}", &proto);

        // Strip path prefix and .proto suffix
        let mod_name = proto
            .strip_prefix("protos/")
            .expect(&format!("Error stripping file {:?}", proto))
            .with_extension("");

        libf.write_all(&format!("pub mod {};\n", mod_name.to_str().unwrap()).as_bytes())
            .expect("Could not write mod in crate's lib file");
    }

    // Add as well the ttrpc file
    libf.write_all(b"pub mod agent_ttrpc;\n")
        .expect("Could not write mod in crate's lib file");

    Codegen::new()
        .out_dir("src")
        .inputs(&protos)
        .include("protos")
        .rust_protobuf()
        .customize(Customize {
            ..Default::default()
        })
        .rust_protobuf_customize(ProtobufCustomize {
            ..Default::default()
        })
        .run()
        .expect("Protocol generation failed.");
}