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

Signed-off-by: Richard Leitner <[email protected]>
hyperium · Jan 17, 2023 · ff642f9 · ff642f9
1 parent 4c69157
commit ff642f9
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 4 deletions.
diff --git a/tests/disable_comments/Cargo.toml b/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" }
diff --git a/tests/extern_path/uuid/Cargo.toml b/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"
diff --git a/tests/wellknown-compiled/Cargo.toml b/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"}
diff --git a/tonic-build/Cargo.toml b/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"
 

diff --git a/tonic-build/src/prost.rs b/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();
         }