feat(build): Builder: add {enum,message}_attributes (#1234)

Signed-off-by: Richard Leitner <[email protected]>

Author: g0hl1n

tests/disable_comments/Cargo.toml

@@ -13,5 +13,5 @@ prost = "0.11"
 tonic = { path = "../../tonic" }
 
 [build-dependencies]
-prost-build = "0.11.4"
+prost-build = "0.11.6"
 tonic-build = { path = "../../tonic-build" }

tests/extern_path/uuid/Cargo.toml

@@ -11,4 +11,4 @@ version = "0.1.0"
 [dependencies]
 prost = "0.11"
 [build-dependencies]
-prost-build = "0.11.4"
+prost-build = "0.11.6"

tests/wellknown-compiled/Cargo.toml

@@ -16,5 +16,5 @@ prost = "0.11"
 tonic = {path = "../../tonic"}
 
 [build-dependencies]
-prost-build = "0.11.4"
+prost-build = "0.11.6"
 tonic-build = {path = "../../tonic-build"}

tonic-build/Cargo.toml

@@ -17,7 +17,7 @@ version = "0.8.4"
 [dependencies]
 prettyplease = { version = "0.1" }
 proc-macro2 = "1.0"
-prost-build = { version = "0.11.4", optional = true }
+prost-build = { version = "0.11.6", optional = true }
 quote = "1.0"
 syn = "1.0"
 

tonic-build/src/prost.rs

@@ -23,6 +23,8 @@ pub fn configure() -> Builder {
         out_dir: None,
         extern_path: Vec::new(),
         field_attributes: Vec::new(),
+        message_attributes: Vec::new(),
+        enum_attributes: Vec::new(),
         type_attributes: Vec::new(),
         server_attributes: Attributes::default(),
         client_attributes: Attributes::default(),
@@ -225,6 +227,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) message_attributes: Vec<(String, String)>,
+    pub(crate) enum_attributes: Vec<(String, String)>,
     pub(crate) server_attributes: Attributes,
     pub(crate) client_attributes: Attributes,
     pub(crate) proto_path: String,
@@ -306,6 +310,28 @@ impl Builder {
         self
     }
 
+    /// Add additional attribute to matched messages.
+    ///
+    /// Passed directly to `prost_build::Config.message_attribute`.
+    pub fn message_attribute<P: AsRef<str>, A: AsRef<str>>(
+        mut self,
+        path: P,
+        attribute: A,
+    ) -> Self {
+        self.message_attributes
+            .push((path.as_ref().to_string(), attribute.as_ref().to_string()));
+        self
+    }
+
+    /// Add additional attribute to matched enums.
+    ///
+    /// Passed directly to `prost_build::Config.enum_attribute`.
+    pub fn enum_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
+        self.enum_attributes
+            .push((path.as_ref().to_string(), attribute.as_ref().to_string()));
+        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,
@@ -447,6 +473,12 @@ impl Builder {
         for (prost_path, attr) in self.type_attributes.iter() {
             config.type_attribute(prost_path, attr);
         }
+        for (prost_path, attr) in self.message_attributes.iter() {
+            config.message_attribute(prost_path, attr);
+        }
+        for (prost_path, attr) in self.enum_attributes.iter() {
+            config.enum_attribute(prost_path, attr);
+        }
         if self.compile_well_known_types {
             config.compile_well_known_types();
         }