feat(build): Add disable_comments option (#1127)

Author: bouzuya

Cargo.toml

@@ -8,6 +8,7 @@ members = [
   "tonic-web", # Non-published crates
   "examples",
   "interop", # Tests
+  "tests/disable_comments",
   "tests/included_service",
   "tests/same_name",
   "tests/service_named_service",

tests/disable_comments/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+authors = ["bouzuya <m@bouzuya.net>"]
+edition = "2018"
+license = "MIT"
+name = "disable-comments"
+publish = false
+version = "0.1.0"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+prost = "0.11"
+tonic = { path = "../../tonic" }
+
+[build-dependencies]
+prost-build = "0.11"
+tonic-build = { path = "../../tonic-build" }

tests/disable_comments/build.rs

@@ -0,0 +1,11 @@
+fn main() {
+    let mut config = prost_build::Config::default();
+    config.disable_comments(&["test.Input1", "test.Output1"]);
+    tonic_build::configure()
+        .disable_comments("test.Service1")
+        .disable_comments("test.Service1.Rpc1")
+        .build_client(true)
+        .build_server(true)
+        .compile_with_config(config, &["proto/test.proto"], &["proto"])
+        .unwrap();
+}

tests/disable_comments/proto/test.proto

@@ -0,0 +1,29 @@
+syntax = "proto3";
+
+package test;
+
+// This comment will be removed.
+service Service1 {
+  // This comment will be removed.
+  rpc Rpc1(Input1) returns (Output1);
+  // This comment will not be removed.
+  rpc Rpc2(Input2) returns (Output2);
+}
+
+// This comment will not be removed.
+service Service2 {
+  // This comment will not be removed.
+  rpc Rpc(Input1) returns (Output1);
+}
+
+// This comment will be removed.
+message Input1 {}
+
+// This comment will not be removed.
+message Input2 {}
+
+// This comment will be removed.
+message Output1 {}
+
+// This comment will not be removed.
+message Output2 {}

tests/disable_comments/src/lib.rs

@@ -0,0 +1,3 @@
+pub mod pb {
+    tonic::include_proto!("test");
+}

tests/disable_comments/tests/disable_comments.rs

@@ -0,0 +1,15 @@
+use std::{fs, path::PathBuf};
+
+#[test]
+fn test() {
+    let path = PathBuf::from(std::env::var("OUT_DIR").unwrap()).join("test.rs");
+    let s = fs::read_to_string(path).unwrap();
+    assert!(!s.contains("This comment will be removed."));
+    let mut count = 0_usize;
+    let mut index = 0_usize;
+    while let Some(found) = s[index..].find("This comment will not be removed.") {
+        index += found + 1;
+        count += 1;
+    }
+    assert_eq!(count, 2 + 3 + 3); // message: 2, client: 3, server: 3
+}

tonic-build/src/client.rs

@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 use super::{Attributes, Method, Service};
 use crate::{generate_doc_comments, naive_snake_case};
 use proc_macro2::TokenStream;
@@ -14,13 +16,19 @@ pub fn generate<T: Service>(
     compile_well_known_types: bool,
     build_transport: bool,
     attributes: &Attributes,
+    disable_comments: &HashSet<String>,
 ) -> 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, emit_package, proto_path, compile_well_known_types);
+    let methods = generate_methods(
+        service,
+        emit_package,
+        proto_path,
+        compile_well_known_types,
+        disable_comments,
+    );
 
     let connect = generate_connect(&service_ident, build_transport);
-    let service_doc = generate_doc_comments(service.comment());
 
     let package = if emit_package { service.package() } else { "" };
     let path = format!(
@@ -30,6 +38,12 @@ pub fn generate<T: Service>(
         service.identifier()
     );
 
+    let service_doc = if disable_comments.contains(&path) {
+        TokenStream::new()
+    } else {
+        generate_doc_comments(service.comment())
+    };
+
     let mod_attributes = attributes.for_mod(package);
     let struct_attributes = attributes.for_struct(&path);
 
@@ -142,6 +156,7 @@ fn generate_methods<T: Service>(
     emit_package: bool,
     proto_path: &str,
     compile_well_known_types: bool,
+    disable_comments: &HashSet<String>,
 ) -> TokenStream {
     let mut stream = TokenStream::new();
     let package = if emit_package { service.package() } else { "" };
@@ -155,7 +170,15 @@ fn generate_methods<T: Service>(
             method.identifier()
         );
 
-        stream.extend(generate_doc_comments(method.comment()));
+        if !disable_comments.contains(&format!(
+            "{}{}{}.{}",
+            package,
+            if package.is_empty() { "" } else { "." },
+            service.identifier(),
+            method.identifier()
+        )) {
+            stream.extend(generate_doc_comments(method.comment()));
+        }
 
         let method = match (method.client_streaming(), method.server_streaming()) {
             (false, false) => generate_unary(method, proto_path, compile_well_known_types, path),

tonic-build/src/manual.rs

@@ -32,6 +32,7 @@ use super::{client, server, Attributes};
 use proc_macro2::TokenStream;
 use quote::ToTokens;
 use std::{
+    collections::HashSet,
     fs,
     path::{Path, PathBuf},
 };
@@ -357,6 +358,7 @@ impl ServiceGenerator {
                 "",    // proto_path, -- not used
                 false, // compile_well_known_types -- not used
                 &Attributes::default(),
+                &HashSet::default(),
             );
             self.servers.extend(server);
         }
@@ -369,6 +371,7 @@ impl ServiceGenerator {
                 false, // compile_well_known_types, -- not used
                 self.builder.build_transport,
                 &Attributes::default(),
+                &HashSet::default(),
             );
             self.clients.extend(client);
         }

tonic-build/src/prost.rs

@@ -3,6 +3,7 @@ use proc_macro2::TokenStream;
 use prost_build::{Config, Method, Service};
 use quote::ToTokens;
 use std::{
+    collections::HashSet,
     ffi::OsString,
     io,
     path::{Path, PathBuf},
@@ -29,6 +30,7 @@ pub fn configure() -> Builder {
         protoc_args: Vec::new(),
         include_file: None,
         emit_rerun_if_changed: std::env::var_os("CARGO").is_some(),
+        disable_comments: HashSet::default(),
     }
 }
 
@@ -163,6 +165,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
                 &self.builder.proto_path,
                 self.builder.compile_well_known_types,
                 &self.builder.server_attributes,
+                &self.builder.disable_comments,
             );
             self.servers.extend(server);
         }
@@ -175,6 +178,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
                 self.builder.compile_well_known_types,
                 self.builder.build_transport,
                 &self.builder.client_attributes,
+                &self.builder.disable_comments,
             );
             self.clients.extend(client);
         }
@@ -229,6 +233,7 @@ pub struct Builder {
     pub(crate) protoc_args: Vec<OsString>,
     pub(crate) include_file: Option<PathBuf>,
     pub(crate) emit_rerun_if_changed: bool,
+    pub(crate) disable_comments: HashSet<String>,
 
     out_dir: Option<PathBuf>,
 }
@@ -354,6 +359,12 @@ impl Builder {
         self
     }
 
+    /// Disable service and rpc comments emission.
+    pub fn disable_comments(mut self, path: impl AsRef<str>) -> Self {
+        self.disable_comments.insert(path.as_ref().to_string());
+        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.

tonic-build/src/server.rs

@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 use super::{Attributes, Method, Service};
 use crate::{generate_doc_comment, generate_doc_comments, naive_snake_case};
 use proc_macro2::{Span, TokenStream};
@@ -14,6 +16,7 @@ pub fn generate<T: Service>(
     proto_path: &str,
     compile_well_known_types: bool,
     attributes: &Attributes,
+    disable_comments: &HashSet<String>,
 ) -> TokenStream {
     let methods = generate_methods(service, proto_path, compile_well_known_types);
 
@@ -22,11 +25,12 @@ pub fn generate<T: Service>(
     let server_mod = quote::format_ident!("{}_server", naive_snake_case(service.name()));
     let generated_trait = generate_trait(
         service,
+        emit_package,
         proto_path,
         compile_well_known_types,
         server_trait.clone(),
+        disable_comments,
     );
-    let service_doc = generate_doc_comments(service.comment());
     let package = if emit_package { service.package() } else { "" };
     // Transport based implementations
     let path = format!(
@@ -35,6 +39,13 @@ pub fn generate<T: Service>(
         if package.is_empty() { "" } else { "." },
         service.identifier()
     );
+
+    let service_doc = if disable_comments.contains(&path) {
+        TokenStream::new()
+    } else {
+        generate_doc_comments(service.comment())
+    };
+
     let named = generate_named(&server_service, &server_trait, &path);
     let mod_attributes = attributes.for_mod(package);
     let struct_attributes = attributes.for_struct(&path);
@@ -167,11 +178,19 @@ pub fn generate<T: Service>(
 
 fn generate_trait<T: Service>(
     service: &T,
+    emit_package: bool,
     proto_path: &str,
     compile_well_known_types: bool,
     server_trait: Ident,
+    disable_comments: &HashSet<String>,
 ) -> TokenStream {
-    let methods = generate_trait_methods(service, proto_path, compile_well_known_types);
+    let methods = generate_trait_methods(
+        service,
+        emit_package,
+        proto_path,
+        compile_well_known_types,
+        disable_comments,
+    );
     let trait_doc = generate_doc_comment(&format!(
         " Generated trait containing gRPC methods that should be implemented for use with {}Server.",
         service.name()
@@ -188,18 +207,31 @@ fn generate_trait<T: Service>(
 
 fn generate_trait_methods<T: Service>(
     service: &T,
+    emit_package: bool,
     proto_path: &str,
     compile_well_known_types: bool,
+    disable_comments: &HashSet<String>,
 ) -> TokenStream {
     let mut stream = TokenStream::new();
 
+    let package = if emit_package { service.package() } else { "" };
     for method in service.methods() {
         let name = quote::format_ident!("{}", method.name());
 
         let (req_message, res_message) =
             method.request_response_name(proto_path, compile_well_known_types);
 
-        let method_doc = generate_doc_comments(method.comment());
+        let method_doc = if disable_comments.contains(&format!(
+            "{}{}{}.{}",
+            package,
+            if package.is_empty() { "" } else { "." },
+            service.identifier(),
+            method.identifier()
+        )) {
+            TokenStream::new()
+        } else {
+            generate_doc_comments(method.comment())
+        };
 
         let method = match (method.client_streaming(), method.server_streaming()) {
             (false, false) => {