From eb0c0c286b08a377e8f7f1a2cf8a55f6aef9caac Mon Sep 17 00:00:00 2001 From: John Doneth Date: Mon, 23 Sep 2019 09:32:44 -0400 Subject: [PATCH] Improve tonic-build configuration (#2) --- tonic-build/Cargo.toml | 1 - tonic-build/src/lib.rs | 170 +++++++++++++++++++++++++++++++--------- tonic-examples/build.rs | 15 +--- tonic-interop/build.rs | 9 +-- 4 files changed, 137 insertions(+), 58 deletions(-) diff --git a/tonic-build/Cargo.toml b/tonic-build/Cargo.toml index b77cd94..20c3292 100644 --- a/tonic-build/Cargo.toml +++ b/tonic-build/Cargo.toml @@ -12,7 +12,6 @@ syn = "1.0" quote = "1.0" proc-macro2 = "1.0" - [features] default = ["transport"] rustfmt = [] diff --git a/tonic-build/src/lib.rs b/tonic-build/src/lib.rs index c35b382..887d469 100644 --- a/tonic-build/src/lib.rs +++ b/tonic-build/src/lib.rs @@ -2,52 +2,124 @@ //! and proto definitiones for use with `tonic`. //! //! # Examples +//! Simple //! //! ```rust,no_run //! fn main() { -//! tonic_build::compile_protos( -//! &["proto/helloworld/helloworld.proto"], -//! &["proto/helloworld"], -//! "helloworld", -//! ) -//! .unwrap(); +//! tonic_build::compile_protos("proto/service.proto").unwrap(); //! } +//! ``` +//! +//! Configuration +//! +//! ```rust,no_run +//! fn main() { +//! tonic_build::configure() +//! .build_server(false) +//! .compile( +//! &["proto/helloworld/helloworld.proto"], +//! &["proto/helloworld"], +//! "helloworld", +//! ) +//! .unwrap(); +//! } +//! ``` use proc_macro2::TokenStream; use prost_build::Config; + #[cfg(feature = "rustfmt")] use std::process::Command; -use std::{io, path, path::Path}; +use std::{ + io, + path::{Path, PathBuf}, +}; mod client; mod service; -pub fn compile_protos

(protos: &[P], includes: &[P], package: &str) -> io::Result<()> -where - P: AsRef, -{ - let out_dir = std::env::var("OUT_DIR").unwrap(); - compile_protos_with_out_dir(protos, includes, package, out_dir.as_str()) +#[derive(Clone)] +pub struct Builder { + build_client: bool, + build_server: bool, + out_dir: PathBuf, } -#[allow(unused_variables)] -pub fn compile_protos_with_out_dir>( - protos: &[P], - includes: &[P], - package: &str, - out_dir: impl AsRef, -) -> io::Result<()> { - let mut config = Config::new(); +impl Builder { + /// Enable or disable gRPC client code generation. + pub fn build_client(mut self, enable: bool) -> Self { + self.build_client = enable; + self + } - config.service_generator(Box::new(ServiceGenerator::default())); - config.out_dir(out_dir.as_ref()); - config.compile_protos(protos, includes)?; + /// Enable or disable gRPC server code generation. + pub fn build_server(mut self, enable: bool) -> Self { + self.build_server = enable; + self + } - #[cfg(feature = "rustfmt")] - fmt( - out_dir.as_ref().to_str().expect("Execpted utf8 out_dir"), - &format!("{}.rs", package), - ); + /// Set the output directory to generate code to. + /// + /// Defaults to the `OUT_DIR` environment variable. + pub fn out_dir(mut self, out_dir: impl AsRef) -> Self { + self.out_dir = out_dir.as_ref().to_path_buf(); + self + } + + /// Compile the .proto files and execute code generation. + #[cfg_attr(not(feature = "rustfmt"), allow(unused_variables))] + pub fn compile>( + self, + protos: &[P], + includes: &[P], + package: &str, + ) -> io::Result<()> { + let mut config = Config::new(); + + config.out_dir(self.out_dir.clone()); + config.service_generator(Box::new(ServiceGenerator::new(self))); + config.compile_protos(protos, includes)?; + + #[cfg(feature = "rustfmt")] + fmt( + out_dir.as_ref().to_str().expect("Execpted utf8 out_dir"), + &format!("{}.rs", package), + ); + + Ok(()) + } +} + +/// Configure tonic-build code generation. +/// +/// Use [`compile_protos`] instead if you don't need to tweak anything. +pub fn configure() -> Builder { + Builder { + build_client: true, + build_server: true, + out_dir: PathBuf::from(std::env::var("OUT_DIR").unwrap()), + } +} + +/// Simple `.proto` compiling. Use [`configure`] instead if you need more options. +/// +/// The include directory will be the parent folder of the specified path. +/// The package name will be the filename without the extension. +pub fn compile_protos(proto_path: impl AsRef) -> io::Result<()> { + let proto_path: &Path = proto_path.as_ref(); + + let package = proto_path + .file_stem() + .expect("file should have a stem if it has an extension") + .to_str() + .expect("expected valid utf-8 filename"); + + // directory the main .proto file resides in + let proto_dir = proto_path + .parent() + .expect("proto file should reside in a directory"); + + self::configure().compile(&[proto_path], &[proto_dir], package)?; Ok(()) } @@ -67,36 +139,58 @@ fn fmt(out_dir: &str, file: &str) { assert!(out.status.success()); } -#[derive(Default)] pub struct ServiceGenerator { + builder: Builder, clients: TokenStream, servers: TokenStream, } +impl ServiceGenerator { + fn new(builder: Builder) -> Self { + ServiceGenerator { + builder, + clients: TokenStream::default(), + servers: TokenStream::default(), + } + } +} + impl prost_build::ServiceGenerator for ServiceGenerator { fn generate(&mut self, service: prost_build::Service, _buf: &mut String) { let path = "super"; - let server = service::generate(&service, path); - self.servers.extend(server); + if self.builder.build_server { + let server = service::generate(&service, path); + self.servers.extend(server); + } - let client = client::generate(&service, path); - self.clients.extend(client); + if self.builder.build_client { + let client = client::generate(&service, path); + self.clients.extend(client); + } } fn finalize(&mut self, buf: &mut String) { - if !self.clients.is_empty() && !self.servers.is_empty() { + if self.builder.build_client && !self.clients.is_empty() { let clients = &self.clients; - let servers = &self.servers; - let service = quote::quote! { + let client_service = quote::quote! { pub mod client { #![allow(unused_variables, dead_code, missing_docs)] use tonic::codegen::*; #clients } + }; + let code = format!("{}", client_service); + buf.push_str(&code); + } + + if self.builder.build_server && !self.servers.is_empty() { + let servers = &self.servers; + + let server_service = quote::quote! { pub mod server { #![allow(unused_variables, dead_code, missing_docs)] use tonic::codegen::*; @@ -105,7 +199,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator { } }; - let code = format!("{}", service); + let code = format!("{}", server_service); buf.push_str(&code); } } diff --git a/tonic-examples/build.rs b/tonic-examples/build.rs index 9b8990f..fcd3a91 100644 --- a/tonic-examples/build.rs +++ b/tonic-examples/build.rs @@ -1,15 +1,4 @@ fn main() { - tonic_build::compile_protos( - &["proto/helloworld/helloworld.proto"], - &["proto/helloworld"], - "helloworld", - ) - .unwrap(); - - tonic_build::compile_protos( - &["proto/routeguide/route_guide.proto"], - &["proto/routeguide"], - "routeguide", - ) - .unwrap(); + tonic_build::compile_protos("proto/helloworld/helloworld.proto").unwrap(); + tonic_build::compile_protos("proto/routeguide/route_guide.proto").unwrap(); } diff --git a/tonic-interop/build.rs b/tonic-interop/build.rs index 9ec7f5d..7783987 100644 --- a/tonic-interop/build.rs +++ b/tonic-interop/build.rs @@ -1,11 +1,8 @@ fn main() { - let files = &["proto/grpc/testing/test.proto"]; - let dirs = &["proto/grpc/testing"]; + let proto = "proto/grpc/testing/test.proto"; - tonic_build::compile_protos(files, dirs, "grpc.testing").unwrap(); + tonic_build::compile_protos(proto).unwrap(); // prevent needing to rebuild if files (or deps) haven't changed - for file in files { - println!("cargo:rerun-if-changed={}", file); - } + println!("cargo:rerun-if-changed={}", proto); }