feat(build): allow to specify trait attributes (#2348)

Author: ya7on

tonic-build/src/lib.rs

@@ -88,6 +88,8 @@ pub struct Attributes {
     module: Vec<(String, String)>,
     /// `struct` attributes.
     structure: Vec<(String, String)>,
+    /// `trait` attributes.
+    trait_attributes: Vec<(String, String)>,
 }
 
 impl Attributes {
@@ -99,6 +101,10 @@ impl Attributes {
         generate_attributes(name, &self.structure)
     }
 
+    fn for_trait(&self, name: &str) -> Vec<syn::Attribute> {
+        generate_attributes(name, &self.trait_attributes)
+    }
+
     /// Add an attribute that will be added to `mod` items matching the given pattern.
     ///
     /// # Examples
@@ -124,6 +130,19 @@ impl Attributes {
     pub fn push_struct(&mut self, pattern: impl Into<String>, attr: impl Into<String>) {
         self.structure.push((pattern.into(), attr.into()));
     }
+
+    /// Add an attribute that will be added to `trait` items matching the given pattern.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use tonic_build::*;
+    /// let mut attributes = Attributes::default();
+    /// attributes.push_trait("Server", "#[mockall::automock]");
+    /// ```
+    pub fn push_trait(&mut self, pattern: impl Into<String>, attr: impl Into<String>) {
+        self.trait_attributes.push((pattern.into(), attr.into()));
+    }
 }
 
 fn format_service_name<T: Service>(service: &T, emit_package: bool) -> String {

tonic-build/src/server.rs

@@ -32,6 +32,7 @@ pub(crate) fn generate_internal<T: Service>(
     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 trait_attributes = attributes.for_trait(service.name());
     let generated_trait = generate_trait(
         service,
         emit_package,
@@ -41,6 +42,7 @@ pub(crate) fn generate_internal<T: Service>(
         disable_comments,
         use_arc_self,
         generate_default_stubs,
+        trait_attributes,
     );
     let package = if emit_package { service.package() } else { "" };
     // Transport based implementations
@@ -203,6 +205,7 @@ fn generate_trait<T: Service>(
     disable_comments: &HashSet<String>,
     use_arc_self: bool,
     generate_default_stubs: bool,
+    trait_attributes: Vec<syn::Attribute>,
 ) -> TokenStream {
     let methods = generate_trait_methods(
         service,
@@ -220,6 +223,7 @@ fn generate_trait<T: Service>(
 
     quote! {
         #trait_doc
+        #(#trait_attributes)*
         #[async_trait]
         pub trait #server_trait : std::marker::Send + std::marker::Sync + 'static {
             #methods

tonic-prost-build/src/lib.rs

@@ -545,6 +545,14 @@ impl Builder {
         self
     }
 
+    /// Add additional attribute to matched traits. Passed directly to
+    /// `prost_build::Config.message_attribute`
+    pub fn trait_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
+        self.server_attributes
+            .push_trait(path.as_ref(), attribute.as_ref());
+        self
+    }
+
     /// Add additional attribute to matched client `mod`s. Passed directly to
     /// `prost_build::Config.message_attribute`
     pub fn client_mod_attribute<P: AsRef<str>, A: AsRef<str>>(