Improve tonic-build configuration (#2)

tonic-build/Cargo.toml

@@ -12,7 +12,6 @@ syn = "1.0"
 quote = "1.0"
 proc-macro2 = "1.0"
 
-
 [features]
 default = ["transport"]
 rustfmt = []

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<P>(protos: &[P], includes: &[P], package: &str) -> io::Result<()>
-where
-    P: AsRef<path::Path>,
-{
-    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<P: AsRef<Path>>(
-    protos: &[P],
-    includes: &[P],
-    package: &str,
-    out_dir: impl AsRef<Path>,
-) -> 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<Path>) -> 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<P: AsRef<Path>>(
+        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<Path>) -> 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);
         }
     }

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();
 }

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);
 }