From 202093c31715b52997c6c206c758924ff5f69bc8 Mon Sep 17 00:00:00 2001
From: Lucio Franco <luciofranco14@gmail.com>
Date: Wed, 1 Apr 2020 11:58:53 -0400
Subject: [PATCH] feat(build): Add support for custom prost config (#318)

Signed-off-by: Lucio Franco <luciofranco14@gmail.com>
---
 tonic-build/src/prost.rs | 126 ++++++++++++++++++++-------------------
 1 file changed, 65 insertions(+), 61 deletions(-)

diff --git a/tonic-build/src/prost.rs b/tonic-build/src/prost.rs
index 5531a9c0b..1351bdbc9 100644
--- a/tonic-build/src/prost.rs
+++ b/tonic-build/src/prost.rs
@@ -5,6 +5,40 @@ use quote::ToTokens;
 use std::io;
 use std::path::{Path, PathBuf};
 
+/// 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: None,
+        extern_path: Vec::new(),
+        field_attributes: Vec::new(),
+        type_attributes: Vec::new(),
+        proto_path: "super".to_string(),
+        #[cfg(feature = "rustfmt")]
+        format: true,
+    }
+}
+
+/// 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: impl AsRef<Path>) -> io::Result<()> {
+    let proto_path: &Path = proto.as_ref();
+
+    // 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])?;
+
+    Ok(())
+}
+
 const PROST_CODEC_PATH: &'static str = "tonic::codec::ProstCodec";
 
 impl crate::Service for Service {
@@ -83,31 +117,6 @@ impl crate::Method for Method {
     }
 }
 
-pub(crate) fn compile<P: AsRef<Path>>(
-    builder: Builder,
-    out_dir: PathBuf,
-    protos: &[P],
-    includes: &[P],
-) -> std::io::Result<()> {
-    let mut config = Config::new();
-
-    config.out_dir(out_dir);
-    for (proto_path, rust_path) in builder.extern_path.iter() {
-        config.extern_path(proto_path, rust_path);
-    }
-    for (prost_path, attr) in builder.field_attributes.iter() {
-        config.field_attribute(prost_path, attr);
-    }
-    for (prost_path, attr) in builder.type_attributes.iter() {
-        config.type_attribute(prost_path, attr);
-    }
-    config.service_generator(Box::new(ServiceGenerator::new(builder)));
-
-    config.compile_protos(protos, includes)?;
-
-    Ok(())
-}
-
 struct ServiceGenerator {
     builder: Builder,
     clients: TokenStream,
@@ -250,7 +259,24 @@ impl Builder {
     }
 
     /// Compile the .proto files and execute code generation.
-    pub fn compile<P: AsRef<Path>>(self, protos: &[P], includes: &[P]) -> io::Result<()> {
+    pub fn compile<P>(self, protos: &[P], includes: &[P]) -> io::Result<()>
+    where
+        P: AsRef<Path>,
+    {
+        self.compile_with_config(Config::new(), protos, includes)
+    }
+
+    /// Compile the .proto files and execute code generation using a
+    /// custom `prost_build::Config`.
+    pub fn compile_with_config<P>(
+        self,
+        mut config: Config,
+        protos: &[P],
+        includes: &[P],
+    ) -> io::Result<()>
+    where
+        P: AsRef<Path>,
+    {
         let out_dir = if let Some(out_dir) = self.out_dir.as_ref() {
             out_dir.clone()
         } else {
@@ -260,7 +286,19 @@ impl Builder {
         #[cfg(feature = "rustfmt")]
         let format = self.format;
 
-        compile(self, out_dir.clone(), protos, includes)?;
+        config.out_dir(out_dir.clone());
+        for (proto_path, rust_path) in self.extern_path.iter() {
+            config.extern_path(proto_path, rust_path);
+        }
+        for (prost_path, attr) in self.field_attributes.iter() {
+            config.field_attribute(prost_path, attr);
+        }
+        for (prost_path, attr) in self.type_attributes.iter() {
+            config.type_attribute(prost_path, attr);
+        }
+        config.service_generator(Box::new(ServiceGenerator::new(self)));
+
+        config.compile_protos(protos, includes)?;
 
         #[cfg(feature = "rustfmt")]
         {
@@ -272,37 +310,3 @@ impl Builder {
         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: None,
-        extern_path: Vec::new(),
-        field_attributes: Vec::new(),
-        type_attributes: Vec::new(),
-        proto_path: "super".to_string(),
-        #[cfg(feature = "rustfmt")]
-        format: true,
-    }
-}
-
-/// 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: impl AsRef<Path>) -> io::Result<()> {
-    let proto_path: &Path = proto.as_ref();
-
-    // 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])?;
-
-    Ok(())
-}