feat(build): Add disable_package_emission option to tonic-build (#556)

Protryon · web-flow · commit 4f5e160679bf · 2021-02-12T10:59:44.000-05:00
diff --git a/tonic-build/src/client.rs b/tonic-build/src/client.rs
@@ -7,10 +7,10 @@ use quote::{format_ident, quote};
 ///
 /// This takes some `Service` and will generate a `TokenStream` that contains
 /// a public module with the generated client.
-pub fn generate<T: Service>(service: &T, proto_path: &str) -> TokenStream {
+pub fn generate<T: Service>(service: &T, emit_package: bool, proto_path: &str) -> TokenStream {
     let service_ident = quote::format_ident!("{}Client", service.name());
     let client_mod = quote::format_ident!("{}_client", naive_snake_case(&service.name()));
-    let methods = generate_methods(service, proto_path);
+    let methods = generate_methods(service, emit_package, proto_path);
 
     let connect = generate_connect(&service_ident);
     let service_doc = generate_doc_comments(service.comment());
@@ -87,18 +87,15 @@ fn generate_connect(_service_ident: &syn::Ident) -> TokenStream {
     TokenStream::new()
 }
 
-fn generate_methods<T: Service>(service: &T, proto_path: &str) -> TokenStream {
+fn generate_methods<T: Service>(service: &T, emit_package: bool, proto_path: &str) -> TokenStream {
     let mut stream = TokenStream::new();
+    let package = if emit_package { service.package() } else { "" };
 
     for method in service.methods() {
         let path = format!(
             "/{}{}{}/{}",
-            service.package(),
-            if service.package().is_empty() {
-                ""
-            } else {
-                "."
-            },
+            package,
+            if package.is_empty() { "" } else { "." },
             service.identifier(),
             method.identifier()
         );
diff --git a/tonic-build/src/prost.rs b/tonic-build/src/prost.rs
@@ -19,6 +19,7 @@ pub fn configure() -> Builder {
         proto_path: "super".to_string(),
         #[cfg(feature = "rustfmt")]
         format: true,
+        emit_package: true,
     }
 }
 
@@ -136,12 +137,20 @@ impl ServiceGenerator {
 impl prost_build::ServiceGenerator for ServiceGenerator {
     fn generate(&mut self, service: prost_build::Service, _buf: &mut String) {
         if self.builder.build_server {
-            let server = server::generate(&service, &self.builder.proto_path);
+            let server = server::generate(
+                &service,
+                self.builder.emit_package,
+                &self.builder.proto_path,
+            );
             self.servers.extend(server);
         }
 
         if self.builder.build_client {
-            let client = client::generate(&service, &self.builder.proto_path);
+            let client = client::generate(
+                &service,
+                self.builder.emit_package,
+                &self.builder.proto_path,
+            );
             self.clients.extend(client);
         }
     }
@@ -184,6 +193,7 @@ pub struct Builder {
     pub(crate) field_attributes: Vec<(String, String)>,
     pub(crate) type_attributes: Vec<(String, String)>,
     pub(crate) proto_path: String,
+    pub(crate) emit_package: bool,
 
     out_dir: Option<PathBuf>,
     #[cfg(feature = "rustfmt")]
@@ -258,6 +268,14 @@ impl Builder {
         self
     }
 
+    /// Emits GRPC endpoints with no attached package. Effectively ignores protofile package declaration from grpc context.
+    ///
+    /// This effectively sets prost's exported package to an empty string.
+    pub fn disable_package_emission(mut self) -> Self {
+        self.emit_package = false;
+        self
+    }
+
     /// Compile the .proto files and execute code generation.
     pub fn compile<P>(self, protos: &[P], includes: &[P]) -> io::Result<()>
     where
diff --git a/tonic-build/src/server.rs b/tonic-build/src/server.rs
@@ -8,24 +8,20 @@ use syn::{Ident, Lit, LitStr};
 ///
 /// This takes some `Service` and will generate a `TokenStream` that contains
 /// a public module containing the server service and handler trait.
-pub fn generate<T: Service>(service: &T, proto_path: &str) -> TokenStream {
+pub fn generate<T: Service>(service: &T, emit_package: bool, proto_path: &str) -> TokenStream {
     let methods = generate_methods(service, proto_path);
 
     let server_service = quote::format_ident!("{}Server", service.name());
     let server_trait = quote::format_ident!("{}", service.name());
     let server_mod = quote::format_ident!("{}_server", naive_snake_case(&service.name()));
     let generated_trait = generate_trait(service, proto_path, server_trait.clone());
     let service_doc = generate_doc_comments(service.comment());
-
+    let package = if emit_package { service.package() } else { "" };
     // Transport based implementations
     let path = format!(
         "{}{}{}",
-        service.package(),
-        if service.package().is_empty() {
-            ""
-        } else {
-            "."
-        },
+        package,
+        if package.is_empty() { "" } else { "." },
         service.identifier()
     );
     let transport = generate_transport(&server_service, &server_trait, &path);
