add config for impl attributes

tonic-build/src/client.rs

@@ -1,5 +1,5 @@
 use super::{Method, Service};
-use crate::{generate_attributes, generate_doc_comments, naive_snake_case};
+use crate::{generate_doc_comments, naive_snake_case, prost::Attributes};
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 
@@ -12,8 +12,7 @@ pub fn generate<T: Service>(
     emit_package: bool,
     proto_path: &str,
     compile_well_known_types: bool,
-    client_mod_attributes: &[(String, String)],
-    client_attributes: &[(String, String)],
+    attributes: &Attributes,
 ) -> TokenStream {
     let service_ident = quote::format_ident!("{}Client", service.name());
     let client_mod = quote::format_ident!("{}_client", naive_snake_case(&service.name()));
@@ -32,24 +31,26 @@ pub fn generate<T: Service>(
         service.identifier()
     );
 
-    let client_mod_attributes = generate_attributes(package, client_mod_attributes);
-    let client_attributes = generate_attributes(&path, client_attributes);
+    let mod_attributes = attributes.for_mod(package);
+    let struct_attributes = attributes.for_struct(&path);
+    let impl_attributes = attributes.for_impl(&path);
 
     quote! {
         /// Generated client implementations.
-        #(#client_mod_attributes)*
+        #(#mod_attributes)*
         pub mod #client_mod {
             #![allow(unused_variables, dead_code, missing_docs)]
             use tonic::codegen::*;
 
             #service_doc
-            #(#client_attributes)*
+            #(#struct_attributes)*
             pub struct #service_ident<T> {
                 inner: tonic::client::Grpc<T>,
             }
 
             #connect
 
+            #(#impl_attributes)*
             impl<T> #service_ident<T>
             where
                 T: tonic::client::GrpcService<tonic::body::BoxBody>,

tonic-build/src/lib.rs

@@ -218,23 +218,6 @@ fn generate_doc_comments<T: AsRef<str>>(comments: &[T]) -> TokenStream {
     stream
 }
 
-// Generates attributes given a list of (`pattern`, `attribute`) pairs. If `pattern` matches `name`, `attribute` will be included.
-pub(crate) 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<_>>()
-}
-
 // Checks whether a path pattern matches a given path.
 pub(crate) fn match_name(pattern: &str, path: &str) -> bool {
     if pattern.is_empty() {

tonic-build/src/prost.rs

@@ -1,3 +1,5 @@
+use crate::match_name;
+
 use super::{client, server};
 use proc_macro2::TokenStream;
 use prost_build::{Config, Method, Service};
@@ -18,10 +20,8 @@ pub fn configure() -> Builder {
         extern_path: Vec::new(),
         field_attributes: Vec::new(),
         type_attributes: Vec::new(),
-        server_mod_attributes: Vec::new(),
-        server_attributes: Vec::new(),
-        client_mod_attributes: Vec::new(),
-        client_attributes: Vec::new(),
+        server_attributes: Attributes::default(),
+        client_attributes: Attributes::default(),
         proto_path: "super".to_string(),
         compile_well_known_types: false,
         #[cfg(feature = "rustfmt")]
@@ -166,8 +166,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
                 self.builder.emit_package,
                 &self.builder.proto_path,
                 self.builder.compile_well_known_types,
-                &self.builder.server_mod_attributes,
-                &self.builder.server_attributes,
+                &self.builder.server_attributes
             );
             self.servers.extend(server);
         }
@@ -178,7 +177,6 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
                 self.builder.emit_package,
                 &self.builder.proto_path,
                 self.builder.compile_well_known_types,
-                &self.builder.client_mod_attributes,
                 &self.builder.client_attributes,
             );
             self.clients.extend(client);
@@ -223,10 +221,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: Vec<(String, String)>,
-    pub(crate) server_mod_attributes: Vec<(String, String)>,
-    pub(crate) client_attributes: Vec<(String, String)>,
-    pub(crate) client_mod_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,
@@ -237,6 +233,7 @@ pub struct Builder {
     format: bool,
 }
 
+
 impl Builder {
     /// Enable or disable gRPC client code generation.
     pub fn build_client(mut self, enable: bool) -> Self {
@@ -309,15 +306,22 @@ impl Builder {
         path: P,
         attribute: A,
     ) -> Self {
-        self.server_mod_attributes
-            .push((path.as_ref().to_string(), attribute.as_ref().to_string()));
+        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((path.as_ref().to_string(), attribute.as_ref().to_string()));
+            .push_struct((path.as_ref().to_string(), attribute.as_ref().to_string()));
+        self
+    }
+
+    /// Add additional attribute to matched service server impls. Matches on the service name.
+    pub fn server_impl_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
+        self.server_attributes
+            .push_impl((path.as_ref().to_string(), attribute.as_ref().to_string()));
         self
     }
 
@@ -327,15 +331,22 @@ impl Builder {
         path: P,
         attribute: A,
     ) -> Self {
-        self.client_mod_attributes
-            .push((path.as_ref().to_string(), attribute.as_ref().to_string()));
+        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((path.as_ref().to_string(), attribute.as_ref().to_string()));
+            .push_struct((path.as_ref().to_string(), attribute.as_ref().to_string()));
+        self
+    }
+
+    /// Add additional attribute to matched service client impls. Matches on the service name.
+    pub fn client_impl_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
+        self.client_attributes
+            .push_impl((path.as_ref().to_string(), attribute.as_ref().to_string()));
         self
     }
 
@@ -436,3 +447,53 @@ impl Builder {
         Ok(())
     }
 }
+
+#[derive(Debug, Default, Clone)]
+pub struct Attributes {
+    pub module: Vec<(String, String)>,
+    pub structure: Vec<(String, String)>,
+    pub implementation: Vec<(String, String)>,
+}
+
+impl Attributes {
+    pub fn for_mod(&self, name: &str) -> Vec<syn::Attribute> {
+        generate_attributes(name, &self.module)
+    }
+
+    pub fn for_struct(&self, name: &str) -> Vec<syn::Attribute> {
+        generate_attributes(name, &self.structure)
+    }
+
+    pub fn for_impl(&self, name: &str) -> Vec<syn::Attribute> {
+        generate_attributes(name, &self.implementation)
+    }
+
+    fn push_mod(&mut self, attr: (String, String)) {
+        self.module.push(attr);
+    }
+
+    fn push_struct(&mut self, attr: (String, String)) {
+        self.structure.push(attr);
+    }
+
+    fn push_impl(&mut self, attr: (String, String)) {
+        self.implementation.push(attr);
+    }
+}
+
+// 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<_>>()
+}

tonic-build/src/server.rs

@@ -1,5 +1,5 @@
 use super::{Method, Service};
-use crate::{generate_attributes, generate_doc_comment, generate_doc_comments, naive_snake_case};
+use crate::{generate_doc_comment, generate_doc_comments, naive_snake_case, prost::Attributes};
 use proc_macro2::{Span, TokenStream};
 use quote::quote;
 use syn::{Attribute, Ident, Lit, LitStr};
@@ -13,8 +13,7 @@ pub fn generate<T: Service>(
     emit_package: bool,
     proto_path: &str,
     compile_well_known_types: bool,
-    server_mod_attributes: &[(String, String)],
-    server_attributes: &[(String, String)],
+    attributes: &Attributes,
 ) -> TokenStream {
     let methods = generate_methods(service, proto_path, compile_well_known_types);
 
@@ -37,20 +36,21 @@ pub fn generate<T: Service>(
         service.identifier()
     );
     let transport = generate_transport(&server_service, &server_trait, &path);
-    let server_mod_attributes = generate_attributes(package, server_mod_attributes);
-    let server_attributes = generate_attributes(&path, server_attributes);
+    let mod_attributes = attributes.for_mod(package);
+    let struct_attributes = attributes.for_struct(&path);
+    let impl_attributes = attributes.for_impl(&path);
 
     quote! {
         /// Generated server implementations.
-        #(#server_mod_attributes)*
+        #(#mod_attributes)*
         pub mod #server_mod {
             #![allow(unused_variables, dead_code, missing_docs)]
             use tonic::codegen::*;
 
             #generated_trait
 
             #service_doc
-            #(#server_attributes)*
+            #(#struct_attributes)*
             #[derive(Debug)]
             pub struct #server_service<T: #server_trait> {
                 inner: _Inner<T>,
@@ -73,6 +73,7 @@ pub fn generate<T: Service>(
                 }
             }
 
+            #(#impl_attributes)*
             impl<T, B> Service<http::Request<B>> for #server_service<T>
                 where
                     T: #server_trait,