feat(build): support adding attributes to clients and servers (#684)

examples/build.rs

@@ -15,6 +15,14 @@ fn main() {
 
     tonic_build::compile_protos("proto/echo/echo.proto").unwrap();
 
+    tonic_build::configure()
+        .server_mod_attribute("attrs", "#[cfg(feature = \"server\")]")
+        .server_attribute("Echo", "#[derive(PartialEq)]")
+        .client_mod_attribute("attrs", "#[cfg(feature = \"client\")]")
+        .client_attribute("Echo", "#[derive(PartialEq)]")
+        .compile(&["proto/attrs/attrs.proto"], &["proto"])
+        .unwrap();
+
     tonic_build::configure()
         .build_server(false)
         .compile(

examples/proto/attrs/attrs.proto

@@ -0,0 +1,19 @@
+syntax = "proto3";
+
+package attrs;
+
+// EchoRequest is the request for echo.
+message EchoRequest {
+  string message = 1;
+}
+
+// EchoResponse is the response for echo.
+message EchoResponse {
+  string message = 1;
+}
+
+// Echo is the echo service.
+service Echo {
+  // UnaryEcho is unary echo.
+  rpc UnaryEcho(EchoRequest) returns (EchoResponse) {}
+}

tonic-build/src/client.rs

@@ -1,4 +1,4 @@
-use super::{Method, Service};
+use super::{Attributes, Method, Service};
 use crate::{generate_doc_comments, naive_snake_case};
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
@@ -12,6 +12,7 @@ pub fn generate<T: Service>(
     emit_package: bool,
     proto_path: &str,
     compile_well_known_types: bool,
+    attributes: &Attributes,
 ) -> TokenStream {
     let service_ident = quote::format_ident!("{}Client", service.name());
     let client_mod = quote::format_ident!("{}_client", naive_snake_case(&service.name()));
@@ -20,13 +21,26 @@ pub fn generate<T: Service>(
     let connect = generate_connect(&service_ident);
     let service_doc = generate_doc_comments(service.comment());
 
+    let package = if emit_package { service.package() } else { "" };
+    let path = format!(
+        "{}{}{}",
+        package,
+        if package.is_empty() { "" } else { "." },
+        service.identifier()
+    );
+
+    let mod_attributes = attributes.for_mod(package);
+    let struct_attributes = attributes.for_struct(&path);
+
     quote! {
         /// Generated client implementations.
+        #(#mod_attributes)*
         pub mod #client_mod {
             #![allow(unused_variables, dead_code, missing_docs)]
             use tonic::codegen::*;
 
             #service_doc
+            #(#struct_attributes)*
             #[derive(Debug, Clone)]
             pub struct #service_ident<T> {
                 inner: tonic::client::Grpc<T>,

tonic-build/src/lib.rs

@@ -155,6 +155,68 @@ pub trait Method {
     ) -> (TokenStream, TokenStream);
 }
 
+/// Attributes that will be added to `mod` and `struct` items.
+#[derive(Debug, Default, Clone)]
+pub struct Attributes {
+    /// `mod` attributes.
+    module: Vec<(String, String)>,
+    /// `struct` attributes.
+    structure: Vec<(String, String)>,
+}
+
+impl Attributes {
+    fn for_mod(&self, name: &str) -> Vec<syn::Attribute> {
+        generate_attributes(name, &self.module)
+    }
+
+    fn for_struct(&self, name: &str) -> Vec<syn::Attribute> {
+        generate_attributes(name, &self.structure)
+    }
+
+    /// Add an attribute that will be added to `mod` items matching the given pattern.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use tonic_build::*;
+    /// let mut attributes = Attributes::default();
+    /// attributes.push_mod("my.proto.package", r#"#[cfg(feature = "server")]"#);
+    /// ```
+    pub fn push_mod(&mut self, pattern: impl Into<String>, attr: impl Into<String>) {
+        self.module.push((pattern.into(), attr.into()));
+    }
+
+    /// Add an attribute that will be added to `struct` items matching the given pattern.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use tonic_build::*;
+    /// let mut attributes = Attributes::default();
+    /// attributes.push_struct("EchoService", "#[derive(PartialEq)]");
+    /// ```
+    pub fn push_struct(&mut self, pattern: impl Into<String>, attr: impl Into<String>) {
+        self.structure.push((pattern.into(), attr.into()));
+    }
+}
+
+// Generates attributes given a list of (`pattern`, `attribute`) pairs. If `pattern` matches `name`, `attribute` will be included.
+fn generate_attributes<'a>(
+    name: &str,
+    attrs: impl IntoIterator<Item = &'a (String, String)>,
+) -> Vec<syn::Attribute> {
+    attrs
+        .into_iter()
+        .filter(|(matcher, _)| match_name(matcher, name))
+        .flat_map(|(_, attr)| {
+            // attributes cannot be parsed directly, so we pretend they're on a struct
+            syn::parse_str::<syn::DeriveInput>(&format!("{}\nstruct fake;", attr))
+                .unwrap()
+                .attrs
+        })
+        .collect::<Vec<_>>()
+}
+
 /// Format files under the out_dir with rustfmt
 #[cfg(feature = "rustfmt")]
 #[cfg_attr(docsrs, doc(cfg(feature = "rustfmt")))]
@@ -218,6 +280,34 @@ fn generate_doc_comments<T: AsRef<str>>(comments: &[T]) -> TokenStream {
     stream
 }
 
+// Checks whether a path pattern matches a given path.
+pub(crate) fn match_name(pattern: &str, path: &str) -> bool {
+    if pattern.is_empty() {
+        false
+    } else if pattern == "." {
+        true
+    } else if pattern == path {
+        true
+    } else {
+        let pattern_segments = pattern.split('.').collect::<Vec<_>>();
+        let path_segments = path.split('.').collect::<Vec<_>>();
+
+        if &pattern[..1] == "." {
+            // prefix match
+            if pattern_segments.len() > path_segments.len() {
+                false
+            } else {
+                pattern_segments[..] == path_segments[..pattern_segments.len()]
+            }
+        // suffix match
+        } else if pattern_segments.len() > path_segments.len() {
+            false
+        } else {
+            pattern_segments[..] == path_segments[path_segments.len() - pattern_segments.len()..]
+        }
+    }
+}
+
 fn naive_snake_case(name: &str) -> String {
     let mut s = String::new();
     let mut it = name.chars().peekable();
@@ -234,14 +324,40 @@ fn naive_snake_case(name: &str) -> String {
     s
 }
 
-#[test]
-fn test_snake_case() {
-    for case in &[
-        ("Service", "service"),
-        ("ThatHasALongName", "that_has_a_long_name"),
-        ("greeter", "greeter"),
-        ("ABCServiceX", "a_b_c_service_x"),
-    ] {
-        assert_eq!(naive_snake_case(case.0), case.1)
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_match_name() {
+        assert!(match_name(".", ".my.protos"));
+        assert!(match_name(".", ".protos"));
+
+        assert!(match_name(".my", ".my"));
+        assert!(match_name(".my", ".my.protos"));
+        assert!(match_name(".my.protos.Service", ".my.protos.Service"));
+
+        assert!(match_name("Service", ".my.protos.Service"));
+
+        assert!(!match_name(".m", ".my.protos"));
+        assert!(!match_name(".p", ".protos"));
+
+        assert!(!match_name(".my", ".myy"));
+        assert!(!match_name(".protos", ".my.protos"));
+        assert!(!match_name(".Service", ".my.protos.Service"));
+
+        assert!(!match_name("service", ".my.protos.Service"));
+    }
+
+    #[test]
+    fn test_snake_case() {
+        for case in &[
+            ("Service", "service"),
+            ("ThatHasALongName", "that_has_a_long_name"),
+            ("greeter", "greeter"),
+            ("ABCServiceX", "a_b_c_service_x"),
+        ] {
+            assert_eq!(naive_snake_case(case.0), case.1)
+        }
     }
 }

tonic-build/src/prost.rs

@@ -1,4 +1,4 @@
-use super::{client, server};
+use super::{client, server, Attributes};
 use proc_macro2::TokenStream;
 use prost_build::{Config, Method, Service};
 use quote::ToTokens;
@@ -18,6 +18,8 @@ pub fn configure() -> Builder {
         extern_path: Vec::new(),
         field_attributes: Vec::new(),
         type_attributes: Vec::new(),
+        server_attributes: Attributes::default(),
+        client_attributes: Attributes::default(),
         proto_path: "super".to_string(),
         compile_well_known_types: false,
         #[cfg(feature = "rustfmt")]
@@ -162,6 +164,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
                 self.builder.emit_package,
                 &self.builder.proto_path,
                 self.builder.compile_well_known_types,
+                &self.builder.server_attributes,
             );
             self.servers.extend(server);
         }
@@ -172,6 +175,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
                 self.builder.emit_package,
                 &self.builder.proto_path,
                 self.builder.compile_well_known_types,
+                &self.builder.client_attributes,
             );
             self.clients.extend(client);
         }
@@ -215,6 +219,8 @@ pub struct Builder {
     pub(crate) extern_path: Vec<(String, String)>,
     pub(crate) field_attributes: Vec<(String, String)>,
     pub(crate) type_attributes: Vec<(String, String)>,
+    pub(crate) server_attributes: Attributes,
+    pub(crate) client_attributes: Attributes,
     pub(crate) proto_path: String,
     pub(crate) emit_package: bool,
     pub(crate) compile_well_known_types: bool,
@@ -291,6 +297,42 @@ impl Builder {
         self
     }
 
+    /// Add additional attribute to matched server `mod`s. Matches on the package name.
+    pub fn server_mod_attribute<P: AsRef<str>, A: AsRef<str>>(
+        mut self,
+        path: P,
+        attribute: A,
+    ) -> Self {
+        self.server_attributes
+            .push_mod(path.as_ref().to_string(), attribute.as_ref().to_string());
+        self
+    }
+
+    /// Add additional attribute to matched service servers. Matches on the service name.
+    pub fn server_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
+        self.server_attributes
+            .push_struct(path.as_ref().to_string(), attribute.as_ref().to_string());
+        self
+    }
+
+    /// Add additional attribute to matched client `mod`s. Matches on the package name.
+    pub fn client_mod_attribute<P: AsRef<str>, A: AsRef<str>>(
+        mut self,
+        path: P,
+        attribute: A,
+    ) -> Self {
+        self.client_attributes
+            .push_mod(path.as_ref().to_string(), attribute.as_ref().to_string());
+        self
+    }
+
+    /// Add additional attribute to matched service clients. Matches on the service name.
+    pub fn client_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
+        self.client_attributes
+            .push_struct(path.as_ref().to_string(), attribute.as_ref().to_string());
+        self
+    }
+
     /// Set the path to where tonic will search for the Request/Response proto structs
     /// live relative to the module where you call `include_proto!`.
     ///

tonic-build/src/server.rs

@@ -1,4 +1,4 @@
-use super::{Method, Service};
+use super::{Attributes, Method, Service};
 use crate::{generate_doc_comment, generate_doc_comments, naive_snake_case};
 use proc_macro2::{Span, TokenStream};
 use quote::quote;
@@ -13,6 +13,7 @@ pub fn generate<T: Service>(
     emit_package: bool,
     proto_path: &str,
     compile_well_known_types: bool,
+    attributes: &Attributes,
 ) -> TokenStream {
     let methods = generate_methods(service, proto_path, compile_well_known_types);
 
@@ -35,6 +36,8 @@ pub fn generate<T: Service>(
         service.identifier()
     );
     let transport = generate_transport(&server_service, &server_trait, &path);
+    let mod_attributes = attributes.for_mod(package);
+    let struct_attributes = attributes.for_struct(&path);
 
     let compression_enabled = cfg!(feature = "compression");
 
@@ -64,13 +67,15 @@ pub fn generate<T: Service>(
 
     quote! {
         /// Generated server implementations.
+        #(#mod_attributes)*
         pub mod #server_mod {
             #![allow(unused_variables, dead_code, missing_docs)]
             use tonic::codegen::*;
 
             #generated_trait
 
             #service_doc
+            #(#struct_attributes)*
             #[derive(Debug)]
             pub struct #server_service<T: #server_trait> {
                 inner: _Inner<T>,