Add concrete example for Config::load_fds()

tokio-rs · May 31, 2024 · cbf0813 · cbf0813
1 parent b952a48
commit cbf0813
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 5 deletions.
diff --git a/prost-build/src/config.rs b/prost-build/src/config.rs
@@ -836,7 +836,7 @@ impl Config {
     }
 
     /// Loads `.proto` files as a [`FileDescriptorSet`]. This allows inspection of the descriptors
-    /// before calling [`Config::compile_fds`]. This could be used to change [`Config`] 
+    /// before calling [`Config::compile_fds`]. This could be used to change [`Config`]
     /// attributes after introspecting what is actually present in the `.proto` files.
     ///
     /// # Example `build.rs`
@@ -848,7 +848,19 @@ impl Config {
     ///   let mut config = Config::new();
     ///   let file_descriptor_set = config.load_fds(&["src/frontend.proto", "src/backend.proto"], &["src"])?;
     ///
-    ///   // Inspect file_descriptor_set and tweak config
+    ///   // Add custom attributes to messages that are service inputs or outputs.
+    ///   for file in &file_descriptor_set.file {
+    ///       for service in &file.service {
+    ///           for method in &service.method {
+    ///               if let Some(input) = &method.input_type {
+    ///                   config.message_attribute(input, "#[derive(custom_proto::Input)]");
+    ///               }
+    ///               if let Some(output) = &method.output_type {
+    ///                   config.message_attribute(output, "#[derive(custom_proto::Output)]");
+    ///               }
+    ///           }
+    ///       }
+    ///   }
     ///
     ///   config.compile_fds(file_descriptor_set)
     /// }

diff --git a/prost-build/src/fixtures/helloworld/_expected_helloworld.rs b/prost-build/src/fixtures/helloworld/_expected_helloworld.rs
@@ -1,12 +1,14 @@
 // This file is @generated by prost-build.
 #[derive(derive_builder::Builder)]
+#[derive(custom_proto::Input)]
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct Message {
     #[prost(string, tag="1")]
     pub say: ::prost::alloc::string::String,
 }
 #[derive(derive_builder::Builder)]
+#[derive(custom_proto::Output)]
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct Response {

diff --git a/prost-build/src/fixtures/helloworld/_expected_helloworld_formatted.rs b/prost-build/src/fixtures/helloworld/_expected_helloworld_formatted.rs
@@ -1,12 +1,14 @@
 // This file is @generated by prost-build.
 #[derive(derive_builder::Builder)]
+#[derive(custom_proto::Input)]
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct Message {
     #[prost(string, tag = "1")]
     pub say: ::prost::alloc::string::String,
 }
 #[derive(derive_builder::Builder)]
+#[derive(custom_proto::Output)]
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct Response {

diff --git a/prost-build/src/lib.rs b/prost-build/src/lib.rs
@@ -390,16 +390,36 @@ mod tests {
         let _ = env_logger::try_init();
         let tempdir = tempfile::tempdir().unwrap();
 
-        Config::new()
+        let mut config = Config::new();
+        config
             .out_dir(tempdir.path())
+            // Add attributes to all messages and enums
             .message_attribute(".", "#[derive(derive_builder::Builder)]")
-            .enum_attribute(".", "#[some_enum_attr(u8)]")
-            .compile_protos(
+            .enum_attribute(".", "#[some_enum_attr(u8)]");
+
+        let fds = config
+            .load_fds(
                 &["src/fixtures/helloworld/hello.proto"],
                 &["src/fixtures/helloworld"],
             )
             .unwrap();
 
+        // Add custom attributes to messages that are service inputs or outputs.
+        for file in &fds.file {
+            for service in &file.service {
+                for method in &service.method {
+                    if let Some(input) = &method.input_type {
+                        config.message_attribute(input, "#[derive(custom_proto::Input)]");
+                    }
+                    if let Some(output) = &method.output_type {
+                        config.message_attribute(output, "#[derive(custom_proto::Output)]");
+                    }
+                }
+            }
+        }
+
+        config.compile_fds(fds).unwrap();
+
         let out_file = tempdir.path().join("helloworld.rs");
         #[cfg(feature = "format")]
         let expected_content =