From f6777849031e1ad9c77d9fc6acb4a572ff138259 Mon Sep 17 00:00:00 2001 From: ziprandom Date: Tue, 1 Aug 2023 11:46:41 +0200 Subject: [PATCH] add Config::preserve_type_name to prevent renaming TODO: unit tests --- prost-build/src/code_generator.rs | 53 +++++++++++++++++++++---------- prost-build/src/lib.rs | 12 +++++++ 2 files changed, 49 insertions(+), 16 deletions(-) diff --git a/prost-build/src/code_generator.rs b/prost-build/src/code_generator.rs index 2a4d24181..08c5cc4f3 100644 --- a/prost-build/src/code_generator.rs +++ b/prost-build/src/code_generator.rs @@ -190,7 +190,8 @@ impl<'a> CodeGenerator<'a> { )); self.push_indent(); self.buf.push_str("pub struct "); - self.buf.push_str(&to_upper_camel(&message_name)); + let type_name = &self.make_type_name(&message_name, &message_name); + self.buf.push_str(&type_name); self.buf.push_str(" {\n"); self.depth += 1; @@ -262,6 +263,19 @@ impl<'a> CodeGenerator<'a> { } } + fn make_type_name(&self, fq_message_name: &str, name: &str) -> String { + if *self + .config + .original_type_names + .get_first(fq_message_name) + .unwrap_or(&false) + { + name.to_string() + } else { + to_upper_camel(&name) + } + } + fn append_type_attributes(&mut self, fq_message_name: &str) { assert_eq!(b'.', fq_message_name.as_bytes()[0]); for attribute in self.config.type_attributes.get(fq_message_name) { @@ -387,7 +401,7 @@ impl<'a> CodeGenerator<'a> { } self.buf.push_str("\\\""); } else if type_ == Type::Enum { - let mut enum_value = to_upper_camel(default); + let mut enum_value = self.make_type_name(&default, &default); if self.config.strip_enum_prefix { // Field types are fully qualified, so we extract // the last segment and strip it from the left @@ -398,7 +412,8 @@ impl<'a> CodeGenerator<'a> { .and_then(|ty| ty.split('.').last()) .unwrap(); - enum_value = strip_enum_prefix(&to_upper_camel(enum_type), &enum_value) + enum_value = + strip_enum_prefix(&self.make_type_name(enum_type, enum_type), &enum_value) } self.buf.push_str(&enum_value); } else { @@ -492,7 +507,7 @@ impl<'a> CodeGenerator<'a> { let name = format!( "{}::{}", to_snake(message_name), - to_upper_camel(oneof.name()) + self.make_type_name(oneof.name(), oneof.name()) ); self.append_doc(fq_message_name, None); self.push_indent(); @@ -538,7 +553,8 @@ impl<'a> CodeGenerator<'a> { )); self.push_indent(); self.buf.push_str("pub enum "); - self.buf.push_str(&to_upper_camel(oneof.name())); + let type_name = &self.make_type_name(oneof.name(), oneof.name()); + self.buf.push_str(&type_name); self.buf.push_str(" {\n"); self.path.push(2); @@ -579,15 +595,14 @@ impl<'a> CodeGenerator<'a> { boxed ); + let type_name = self.make_type_name(field.name(), field.name()); if boxed { self.buf.push_str(&format!( "{}(::prost::alloc::boxed::Box<{}>),\n", - to_upper_camel(field.name()), - ty + type_name, ty )); } else { - self.buf - .push_str(&format!("{}({}),\n", to_upper_camel(field.name()), ty)); + self.buf.push_str(&format!("{}({}),\n", type_name, ty)); } } self.depth -= 1; @@ -626,7 +641,7 @@ impl<'a> CodeGenerator<'a> { debug!(" enum: {:?}", desc.name()); let proto_enum_name = desc.name(); - let enum_name = to_upper_camel(proto_enum_name); + let enum_name = self.make_type_name(proto_enum_name, proto_enum_name); let enum_values = &desc.value; let fq_proto_enum_name = format!( @@ -657,8 +672,12 @@ impl<'a> CodeGenerator<'a> { self.buf.push_str(&enum_name); self.buf.push_str(" {\n"); - let variant_mappings = - build_enum_value_mappings(&enum_name, self.config.strip_enum_prefix, enum_values); + let variant_mappings = build_enum_value_mappings( + &enum_name, + self.config.strip_enum_prefix, + enum_values, + |field_name: &str| self.make_type_name(&enum_name, field_name), + ); self.depth += 1; self.path.push(2); @@ -817,7 +836,7 @@ impl<'a> CodeGenerator<'a> { self.path.pop(); let service = Service { - name: to_upper_camel(&name), + name: self.make_type_name(&name, &name), proto_name: name, package: self.package.clone(), comments, @@ -861,7 +880,7 @@ impl<'a> CodeGenerator<'a> { self.buf.push_str("}\n"); } - fn resolve_type(&self, field: &FieldDescriptorProto, fq_message_name: &str) -> String { + fn resolve_type(&mut self, field: &FieldDescriptorProto, fq_message_name: &str) -> String { let prost_path = self.config.prost_path.as_deref().unwrap_or("::prost"); match field.r#type() { @@ -915,7 +934,7 @@ impl<'a> CodeGenerator<'a> { local_path .map(|_| "super".to_string()) .chain(ident_path.map(to_snake)) - .chain(iter::once(to_upper_camel(ident_type))) + .chain(iter::once(self.make_type_name(ident_type, ident_type))) .join("::") } @@ -1142,6 +1161,7 @@ fn build_enum_value_mappings<'a>( generated_enum_name: &str, do_strip_enum_prefix: bool, enum_values: &'a [EnumValueDescriptorProto], + mut variant_name_cb: impl FnMut(&str) -> String, ) -> Vec> { let mut numbers = HashSet::new(); let mut generated_names = HashMap::new(); @@ -1154,7 +1174,8 @@ fn build_enum_value_mappings<'a>( continue; } - let mut generated_variant_name = to_upper_camel(value.name()); + let mut generated_variant_name = variant_name_cb(value.name()); + if do_strip_enum_prefix { generated_variant_name = strip_enum_prefix(generated_enum_name, &generated_variant_name); diff --git a/prost-build/src/lib.rs b/prost-build/src/lib.rs index 80179c557..bcbc24249 100644 --- a/prost-build/src/lib.rs +++ b/prost-build/src/lib.rs @@ -243,6 +243,7 @@ pub struct Config { service_generator: Option>, map_type: PathMap, bytes_type: PathMap, + original_type_names: PathMap, type_attributes: PathMap, message_attributes: PathMap, enum_attributes: PathMap, @@ -471,6 +472,16 @@ impl Config { self } + // Don't rename the type at the given path + pub fn preserve_type_name

(&mut self, path: P) -> &mut Self + where + P: AsRef, + { + self.original_type_names + .insert(path.as_ref().to_string(), true); + self + } + /// Add additional attribute to matched messages. /// /// # Arguments @@ -1233,6 +1244,7 @@ impl default::Default for Config { service_generator: None, map_type: PathMap::default(), bytes_type: PathMap::default(), + original_type_names: PathMap::default(), type_attributes: PathMap::default(), message_attributes: PathMap::default(), enum_attributes: PathMap::default(),