feat(build): Allow generated Message derive path to be changed (#729)

Co-authored-by: Lucio Franco <[email protected]>
tokio-rs · Oct 21, 2022 · 73150d8 · 73150d8
1 parent ec385dc
commit 73150d8
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 5 deletions.
diff --git a/prost-build/src/code_generator.rs b/prost-build/src/code_generator.rs
@@ -183,8 +183,10 @@ impl<'a> CodeGenerator<'a> {
         self.append_doc(&fq_message_name, None);
         self.append_type_attributes(&fq_message_name);
         self.push_indent();
-        self.buf
-            .push_str("#[derive(Clone, PartialEq, ::prost::Message)]\n");
+        self.buf.push_str(&format!(
+            "#[derive(Clone, PartialEq, {}::Message)]\n",
+            self.config.prost_path.as_deref().unwrap_or("::prost")
+        ));
         self.push_indent();
         self.buf.push_str("pub struct ");
         self.buf.push_str(&to_upper_camel(&message_name));
@@ -498,8 +500,10 @@ impl<'a> CodeGenerator<'a> {
         let oneof_name = format!("{}.{}", fq_message_name, oneof.name());
         self.append_type_attributes(&oneof_name);
         self.push_indent();
-        self.buf
-            .push_str("#[derive(Clone, PartialEq, ::prost::Oneof)]\n");
+        self.buf.push_str(&format!(
+            "#[derive(Clone, PartialEq, {}::Oneof)]\n",
+            self.config.prost_path.as_deref().unwrap_or("::prost")
+        ));
         self.push_indent();
         self.buf.push_str("pub enum ");
         self.buf.push_str(&to_upper_camel(oneof.name()));
@@ -605,7 +609,7 @@ impl<'a> CodeGenerator<'a> {
         self.append_type_attributes(&fq_proto_enum_name);
         self.push_indent();
         self.buf.push_str(
-            "#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]\n",
+            &format!("#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, {}::Enumeration)]\n",self.config.prost_path.as_deref().unwrap_or("::prost")),
         );
         self.push_indent();
         self.buf.push_str("#[repr(i32)]\n");

diff --git a/prost-build/src/lib.rs b/prost-build/src/lib.rs
@@ -253,6 +253,7 @@ pub struct Config {
     disable_comments: PathMap<()>,
     skip_protoc_run: bool,
     include_file: Option<PathBuf>,
+    prost_path: Option<String>,
 }
 
 impl Config {
@@ -707,6 +708,17 @@ impl Config {
         self
     }
 
+    /// Configures the path that's used for deriving `Message` for generated messages.
+    /// This is mainly useful for generating crates that wish to re-export prost.
+    /// Defaults to `::prost::Message` if not specified.
+    pub fn prost_path<S>(&mut self, path: S) -> &mut Self
+    where
+        S: Into<String>,
+    {
+        self.prost_path = Some(path.into());
+        self
+    }
+
     /// Add an argument to the `protoc` protobuf compilation invocation.
     ///
     /// # Example `build.rs`
@@ -1068,6 +1080,7 @@ impl default::Default for Config {
             disable_comments: PathMap::default(),
             skip_protoc_run: false,
             include_file: None,
+            prost_path: None,
         }
     }
 }
@@ -1088,6 +1101,7 @@ impl fmt::Debug for Config {
             .field("default_package_filename", &self.default_package_filename)
             .field("protoc_args", &self.protoc_args)
             .field("disable_comments", &self.disable_comments)
+            .field("prost_path", &self.prost_path)
             .finish()
     }
 }