Backport CustomizedCallback features from protobuf v3

protobuf-codegen/src/customize/customize_callback.rs

@@ -0,0 +1,47 @@
+use protobuf::descriptor::DescriptorProto;
+use protobuf::descriptor::EnumDescriptorProto;
+use protobuf::descriptor::FieldDescriptorProto;
+use protobuf::descriptor::OneofDescriptorProto;
+
+use crate::code_writer::CodeWriter;
+
+use super::Customize;
+
+/// Dynamic callback to customize code generation.
+pub trait CustomizeCallback: 'static {
+    fn message(&self, message: &DescriptorProto) -> Customize {
+        let _ = message;
+        Customize::default()
+    }
+
+    fn field(&self, field: &FieldDescriptorProto) -> Customize {
+        let _ = field;
+        Customize::default()
+    }
+
+    fn special_field(&self, message: &DescriptorProto, field: &str) -> Customize {
+        let _ = (message, field);
+        Customize::default()
+    }
+
+    fn enumeration(&self, enum_type: &EnumDescriptorProto) -> Customize {
+        let _ = enum_type;
+        Customize::default()
+    }
+
+    fn oneof(&self, oneof: &OneofDescriptorProto) -> Customize {
+        let _ = oneof;
+        Customize::default()
+    }
+}
+
+pub(crate) struct CustomizeCallbackDefault;
+
+impl CustomizeCallback for CustomizeCallbackDefault {}
+
+/// Write output that generated by [`CustomizeCallback`].
+pub fn write_customize_callback(w: &mut CodeWriter, customize: &Customize) {
+    for line in customize.before.iter().flat_map(|s| s.lines()) {
+        w.write_line(line);
+    }
+}

protobuf-codegen/src/customize.rs → protobuf-codegen/src/customize/mod.rs

@@ -3,9 +3,16 @@ use protobuf::descriptor::FileOptions;
 use protobuf::descriptor::MessageOptions;
 use protobuf::rustproto;
 
+mod customize_callback;
+pub use self::customize_callback::write_customize_callback;
+pub use self::customize_callback::CustomizeCallback;
+pub(crate) use self::customize_callback::CustomizeCallbackDefault;
+
 /// Specifies style of generated code.
 #[derive(Default, Debug, Clone)]
 pub struct Customize {
+    /// Code to insert before the element in the generated file.
+    pub before: Option<String>,
     /// Make oneof enum public.
     pub expose_oneof: Option<bool>,
     /// When true all fields are public, and accessors are not generated
@@ -42,8 +49,24 @@ pub enum CustomizeParseParameterError {
 pub type CustomizeParseParameterResult<T> = Result<T, CustomizeParseParameterError>;
 
 impl Customize {
+    /// Insert code before the element in the generated file
+    /// (e. g. serde annotations, see
+    /// [example here](https://github.com/stepancheg/rust-protobuf/tree/master/protobuf-examples/customize-serde)).
+    pub fn before(mut self, before: &str) -> Self {
+        self.before = Some(before.to_owned());
+        self
+    }
+
+    /// Update fields of self with fields defined in other customize
+    pub fn update_from_callback(&mut self, that: &Customize) {
+        self.before = that.before.clone();
+    }
+
     /// Update fields of self with fields defined in other customize
     pub fn update_with(&mut self, that: &Customize) {
+        if let Some(v) = &that.before {
+            self.before = Some(v.clone());
+        }
         if let Some(v) = that.expose_oneof {
             self.expose_oneof = Some(v);
         }
@@ -126,6 +149,7 @@ impl Customize {
 }
 
 pub fn customize_from_rustproto_for_message(source: &MessageOptions) -> Customize {
+    let before = None;
     let expose_oneof = rustproto::exts::expose_oneof.get(source);
     let expose_fields = rustproto::exts::expose_fields.get(source);
     let generate_accessors = rustproto::exts::generate_accessors.get(source);
@@ -136,6 +160,7 @@ pub fn customize_from_rustproto_for_message(source: &MessageOptions) -> Customiz
     let lite_runtime = None;
     let inside_protobuf = None;
     Customize {
+        before,
         expose_oneof,
         expose_fields,
         generate_accessors,
@@ -150,6 +175,7 @@ pub fn customize_from_rustproto_for_message(source: &MessageOptions) -> Customiz
 }
 
 pub fn customize_from_rustproto_for_field(source: &FieldOptions) -> Customize {
+    let before = None;
     let expose_oneof = None;
     let expose_fields = rustproto::exts::expose_fields_field.get(source);
     let generate_accessors = rustproto::exts::generate_accessors_field.get(source);
@@ -161,6 +187,7 @@ pub fn customize_from_rustproto_for_field(source: &FieldOptions) -> Customize {
     let lite_runtime = None;
     let inside_protobuf = None;
     Customize {
+        before,
         expose_oneof,
         expose_fields,
         generate_accessors,
@@ -175,6 +202,7 @@ pub fn customize_from_rustproto_for_field(source: &FieldOptions) -> Customize {
 }
 
 pub fn customize_from_rustproto_for_file(source: &FileOptions) -> Customize {
+    let before = None;
     let expose_oneof = rustproto::exts::expose_oneof_all.get(source);
     let expose_fields = rustproto::exts::expose_fields_all.get(source);
     let generate_accessors = rustproto::exts::generate_accessors_all.get(source);
@@ -185,6 +213,7 @@ pub fn customize_from_rustproto_for_file(source: &FileOptions) -> Customize {
     let lite_runtime = rustproto::exts::lite_runtime_all.get(source);
     let inside_protobuf = None;
     Customize {
+        before,
         expose_oneof,
         expose_fields,
         generate_accessors,

protobuf-codegen/src/enums.rs

@@ -3,6 +3,8 @@ use std::collections::HashSet;
 use protobuf::descriptor::*;
 use protobuf::descriptorx::*;
 
+use crate::customize::CustomizeCallback;
+
 use super::code_writer::*;
 use super::customize::Customize;
 use rust_types_values::type_name_to_rust_relative;
@@ -53,6 +55,7 @@ impl<'a> EnumGen<'a> {
         enum_with_scope: &'a EnumWithScope<'a>,
         current_file: &FileDescriptorProto,
         customize: &Customize,
+        customize_callback: &'a dyn CustomizeCallback,
         root_scope: &RootScope,
     ) -> EnumGen<'a> {
         let rust_name = if enum_with_scope.get_scope().get_file_descriptor().get_name()
@@ -76,12 +79,14 @@ impl<'a> EnumGen<'a> {
                 .get_optimize_for()
                 == FileOptions_OptimizeMode::LITE_RUNTIME
         });
+        let mut customize = customize.clone();
+        customize.update_from_callback(&customize_callback.enumeration(enum_with_scope.en));
 
         EnumGen {
             enum_with_scope,
             type_name: rust_name,
             lite_runtime: lite_runtime,
-            customize: customize.clone(),
+            customize,
         }
     }
 
@@ -164,6 +169,7 @@ impl<'a> EnumGen<'a> {
         }
         w.derive(&derive);
         serde::write_serde_attr(w, &self.customize, "derive(Serialize, Deserialize)");
+        crate::customize::write_customize_callback(w, &self.customize);
         let ref type_name = self.type_name;
         w.expr_block(&format!("pub enum {}", type_name), |w| {
             for value in self.values_all() {
@@ -243,7 +249,7 @@ impl<'a> EnumGen<'a> {
     fn write_impl_value(&self, w: &mut CodeWriter) {
         w.impl_for_block("::protobuf::reflect::ProtobufValue", &self.type_name, |w| {
             w.def_fn(
-                "as_ref(&self) -> ::protobuf::reflect::ProtobufValueRef",
+                "as_ref(&self) -> ::protobuf::reflect::ProtobufValueRef<'_>",
                 |w| w.write_line("::protobuf::reflect::ProtobufValueRef::Enum(self.descriptor())"),
             )
         })

protobuf-codegen/src/field.rs

@@ -5,6 +5,8 @@ use protobuf::rust;
 use protobuf::text_format;
 use protobuf::wire_format;
 
+use crate::customize::CustomizeCallback;
+
 use super::code_writer::CodeWriter;
 use super::enums::*;
 use super::rust_types_values::*;
@@ -290,6 +292,7 @@ fn field_elem(
     root_scope: &RootScope,
     parse_map: bool,
     customize: &Customize,
+    customize_callback: &dyn CustomizeCallback,
 ) -> (FieldElem, Option<EnumValueGen>) {
     if field.field.get_field_type() == FieldDescriptorProto_Type::TYPE_GROUP {
         (FieldElem::Group, None)
@@ -316,8 +319,8 @@ fn field_elem(
                     (parse_map, message_with_scope.map_entry())
                 {
                     Some(Box::new(EntryKeyValue(
-                        field_elem(&key, root_scope, false, customize).0,
-                        field_elem(&value, root_scope, false, customize).0,
+                        field_elem(&key, root_scope, false, customize, customize_callback).0,
+                        field_elem(&value, root_scope, false, customize, customize_callback).0,
                     )))
                 } else {
                     None
@@ -335,6 +338,7 @@ fn field_elem(
                     &enum_with_scope,
                     field.message.get_scope().get_file_descriptor(),
                     customize,
+                    customize_callback,
                     root_scope,
                 );
                 let ev = if field.field.has_default_value() {
@@ -426,13 +430,16 @@ impl<'a> FieldGen<'a> {
         field: FieldWithContext<'a>,
         root_scope: &'a RootScope<'a>,
         customize: &Customize,
+        customize_callback: &'a dyn CustomizeCallback,
     ) -> FieldGen<'a> {
         let mut customize = customize.clone();
         customize.update_with(&customize_from_rustproto_for_field(
             &field.field.get_options(),
         ));
+        customize.update_from_callback(&customize_callback.field(field.field));
 
-        let (elem, enum_default_value) = field_elem(&field, root_scope, true, &customize);
+        let (elem, enum_default_value) =
+            field_elem(&field, root_scope, true, &customize, customize_callback);
 
         let generate_accessors = customize.generate_accessors.unwrap_or(true);
 
@@ -486,6 +493,10 @@ impl<'a> FieldGen<'a> {
         }
     }
 
+    pub(crate) fn write_customize_callback(&self, w: &mut CodeWriter) {
+        crate::customize::write_customize_callback(w, &self.customize);
+    }
+
     fn tag_size(&self) -> u32 {
         rt::tag_size(self.proto_field.number())
     }

protobuf-codegen/src/lib.rs

@@ -1,5 +1,5 @@
-extern crate protobuf;
 extern crate heck;
+extern crate protobuf;
 
 use std::collections::hash_map::HashMap;
 use std::fmt::Write as FmtWrite;
@@ -8,6 +8,7 @@ use std::io;
 use std::io::Write;
 use std::path::Path;
 
+use customize::CustomizeCallbackDefault;
 use protobuf::compiler_plugin;
 use protobuf::descriptor::*;
 use protobuf::descriptorx::*;
@@ -29,7 +30,7 @@ mod serde;
 mod well_known_types;
 
 use customize::customize_from_rustproto_for_file;
-pub use customize::Customize;
+pub use customize::{Customize, CustomizeCallback};
 
 pub mod code_writer;
 
@@ -126,6 +127,7 @@ fn gen_file(
     _files_map: &HashMap<&str, &FileDescriptorProto>,
     root_scope: &RootScope,
     customize: &Customize,
+    customize_callback: &dyn CustomizeCallback,
 ) -> Option<compiler_plugin::GenResult> {
     // TODO: use it
     let mut customize = customize.clone();
@@ -166,12 +168,12 @@ fn gen_file(
             // ignore map entries, because they are not used in map fields
             if message.map_entry().is_none() {
                 w.write_line("");
-                MessageGen::new(message, &root_scope, &customize).write(&mut w);
+                MessageGen::new(message, &root_scope, &customize, customize_callback).write(&mut w);
             }
         }
         for enum_type in &scope.get_enums() {
             w.write_line("");
-            EnumGen::new(enum_type, file, &customize, root_scope).write(&mut w);
+            EnumGen::new(enum_type, file, &customize, customize_callback, root_scope).write(&mut w);
         }
 
         write_extensions(file, &root_scope, &mut w, &customize);
@@ -195,6 +197,20 @@ pub fn gen(
     file_descriptors: &[FileDescriptorProto],
     files_to_generate: &[String],
     customize: &Customize,
+) -> Vec<compiler_plugin::GenResult> {
+    gen_with_callback(
+        file_descriptors,
+        files_to_generate,
+        customize,
+        &CustomizeCallbackDefault,
+    )
+}
+
+pub fn gen_with_callback(
+    file_descriptors: &[FileDescriptorProto],
+    files_to_generate: &[String],
+    customize: &Customize,
+    customize_callback: &dyn CustomizeCallback,
 ) -> Vec<compiler_plugin::GenResult> {
     let root_scope = RootScope {
         file_descriptors: file_descriptors,
@@ -211,7 +227,13 @@ pub fn gen(
             "file not found in file descriptors: {:?}, files: {:?}",
             file_name, all_file_names
         ));
-        results.extend(gen_file(file, &files_map, &root_scope, customize));
+        results.extend(gen_file(
+            file,
+            &files_map,
+            &root_scope,
+            customize,
+            customize_callback,
+        ));
     }
     results
 }
@@ -222,7 +244,28 @@ pub fn gen_and_write(
     out_dir: &Path,
     customize: &Customize,
 ) -> io::Result<()> {
-    let results = gen(file_descriptors, files_to_generate, customize);
+    gen_and_write_with_callback(
+        file_descriptors,
+        files_to_generate,
+        out_dir,
+        customize,
+        &CustomizeCallbackDefault,
+    )
+}
+
+pub fn gen_and_write_with_callback(
+    file_descriptors: &[FileDescriptorProto],
+    files_to_generate: &[String],
+    out_dir: &Path,
+    customize: &Customize,
+    customize_callback: &dyn CustomizeCallback,
+) -> io::Result<()> {
+    let results = gen_with_callback(
+        file_descriptors,
+        files_to_generate,
+        customize,
+        customize_callback,
+    );
 
     for r in &results {
         let mut file_path = out_dir.to_owned();