Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 21 additions & 3 deletions pbjson-build/src/generator/enumeration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,9 @@ use super::{
use crate::descriptor::{EnumDescriptor, TypePath};
use crate::generator::write_fields_array;
use crate::resolver::Resolver;
use heck::ToUpperCamelCase;
use std::collections::HashSet;
use itertools::Itertools;
use std::io::{Result, Write};

pub fn generate_enum<W: Write>(
Expand All @@ -20,6 +22,7 @@ pub fn generate_enum<W: Write>(
descriptor: &EnumDescriptor,
writer: &mut W,
use_integers_for_enums: bool,
support_camel_case_for_enum_deserialization: bool,
) -> Result<()> {
let rust_type = resolver.rust_type(path);

Expand Down Expand Up @@ -74,7 +77,13 @@ pub fn generate_enum<W: Write>(
// Generate Deserialize
write_deserialize_start(0, &rust_type, writer)?;
write_fields_array(writer, 2, variants.iter().map(|(name, _, _)| name.as_str()))?;
write_visitor(writer, 2, &rust_type, &variants)?;
write_visitor(
writer,
2,
&rust_type,
&variants,
support_camel_case_for_enum_deserialization,
)?;

// Use deserialize_any to allow users to provide integers or strings
writeln!(
Expand All @@ -92,6 +101,7 @@ fn write_visitor<W: Write>(
indent: usize,
rust_type: &str,
variants: &[(String, i32, String)],
support_camel_case_for_enum_deserialization: bool,
) -> Result<()> {
// Protobuf supports deserialization of enumerations both from string and integer values
writeln!(
Expand Down Expand Up @@ -139,14 +149,22 @@ fn write_visitor<W: Write>(

writeln!(writer, "{}match value {{", Indent(indent + 2))?;
for (variant_name, _, rust_variant) in variants {
let mut variants = vec![variant_name.to_string()];
if support_camel_case_for_enum_deserialization {
let camel_case = variant_name.to_upper_camel_case();
variants.push(camel_case);
variants.dedup();
}
let variants = variants.into_iter().map(|variant| format!("\"{variant}\"")).join(" | ");
writeln!(
writer,
"{}\"{}\" => Ok({}::{}),",
"{}{} => Ok({}::{}),",
Indent(indent + 3),
variant_name,
variants,
rust_type,
rust_variant
)?;

}

writeln!(
Expand Down
9 changes: 9 additions & 0 deletions pbjson-build/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,7 @@ pub struct Builder {
btree_map_paths: Vec<String>,
emit_fields: bool,
use_integers_for_enums: bool,
support_camel_case_for_enum_deserialization: bool,
preserve_proto_field_names: bool,
}

Expand Down Expand Up @@ -186,12 +187,19 @@ impl Builder {
self.emit_fields = true;
self
}

// print integers instead of enum names.
pub fn use_integers_for_enums(&mut self) -> &mut Self {
self.use_integers_for_enums = true;
self
}

// Deserialize camelCased strings as enums.
pub fn support_camel_case_for_enum_deserialization(&mut self) -> &mut Self {
self.support_camel_case_for_enum_deserialization = true;
self
}

/// Output fields with their original names as defined in their proto schemas, instead of
/// lowerCamelCase
pub fn preserve_proto_field_names(&mut self) -> &mut Self {
Expand Down Expand Up @@ -277,6 +285,7 @@ impl Builder {
descriptor,
writer,
self.use_integers_for_enums,
self.support_camel_case_for_enum_deserialization,
)?,
Descriptor::Message(descriptor) => {
if let Some(message) = resolve_message(&self.descriptors, descriptor) {
Expand Down