Merge pull request #69 from lebai-robotics/main

feat: add flag to print integers instead of enum names

.circleci/config.yml

@@ -134,6 +134,9 @@ jobs:
       - run:
           name: Cargo test (emit fields)
           command: cargo test --workspace --features emit-fields
+      - run:
+          name: Cargo test (use integers for enums)
+          command: cargo test --workspace --features use-integers-for-enums
       - cache_save
 
 workflows:

pbjson-build/src/generator/enumeration.rs

@@ -18,6 +18,7 @@ pub fn generate_enum<W: Write>(
     path: &TypePath,
     descriptor: &EnumDescriptor,
     writer: &mut W,
+    use_integers_for_enums: bool,
 ) -> Result<()> {
     let rust_type = resolver.rust_type(path);
 
@@ -26,31 +27,48 @@ pub fn generate_enum<W: Write>(
         .iter()
         .map(|variant| {
             let variant_name = variant.name.clone().unwrap();
+            let variant_number = variant.number();
             let rust_variant = resolver.rust_variant(path, &variant_name);
-            (variant_name, rust_variant)
+            (variant_name, variant_number, rust_variant)
         })
         .collect();
 
     // Generate Serialize
     write_serialize_start(0, &rust_type, writer)?;
-    writeln!(writer, "{}let variant = match self {{", Indent(2))?;
-    for (variant_name, rust_variant) in &variants {
-        writeln!(
-            writer,
-            "{}Self::{} => \"{}\",",
-            Indent(3),
-            rust_variant,
-            variant_name
-        )?;
+    if use_integers_for_enums {
+        writeln!(writer, "{}let variant = match self {{", Indent(2))?;
+        for (_, variant_number, rust_variant) in &variants {
+            writeln!(
+                writer,
+                "{}Self::{} => {},",
+                Indent(3),
+                rust_variant,
+                variant_number
+            )?;
+        }
+        writeln!(writer, "{}}};", Indent(2))?;
+
+        writeln!(writer, "{}serializer.serialize_i32(variant)", Indent(2))?;
+    } else {
+        writeln!(writer, "{}let variant = match self {{", Indent(2))?;
+        for (variant_name, _, rust_variant) in &variants {
+            writeln!(
+                writer,
+                "{}Self::{} => \"{}\",",
+                Indent(3),
+                rust_variant,
+                variant_name
+            )?;
+        }
+        writeln!(writer, "{}}};", Indent(2))?;
+
+        writeln!(writer, "{}serializer.serialize_str(variant)", Indent(2))?;
     }
-    writeln!(writer, "{}}};", Indent(2))?;
-
-    writeln!(writer, "{}serializer.serialize_str(variant)", Indent(2))?;
     write_serialize_end(0, writer)?;
 
     // Generate Deserialize
     write_deserialize_start(0, &rust_type, writer)?;
-    write_fields_array(writer, 2, variants.iter().map(|(name, _)| name.as_str()))?;
+    write_fields_array(writer, 2, variants.iter().map(|(name, _, _)| name.as_str()))?;
     write_visitor(writer, 2, &rust_type, &variants)?;
 
     // Use deserialize_any to allow users to provide integers or strings
@@ -68,7 +86,7 @@ fn write_visitor<W: Write>(
     writer: &mut W,
     indent: usize,
     rust_type: &str,
-    variants: &[(String, String)],
+    variants: &[(String, i32, String)],
 ) -> Result<()> {
     // Protobuf supports deserialization of enumerations both from string and integer values
     writeln!(
@@ -117,7 +135,7 @@ fn write_visitor<W: Write>(
     )?;
 
     writeln!(writer, "{}match value {{", Indent(indent + 2))?;
-    for (variant_name, rust_variant) in variants {
+    for (variant_name, _, rust_variant) in variants {
         writeln!(
             writer,
             "{}\"{}\" => Ok({}::{}),",

pbjson-build/src/lib.rs

@@ -106,6 +106,7 @@ pub struct Builder {
     ignore_unknown_fields: bool,
     btree_map_paths: Vec<String>,
     emit_fields: bool,
+    use_integers_for_enums: bool,
 }
 
 impl Builder {
@@ -184,6 +185,11 @@ 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
+    }
 
     /// Generates code for all registered types where `prefixes` contains a prefix of
     /// the fully-qualified path of the type
@@ -257,9 +263,13 @@ impl Builder {
             );
 
             match descriptor {
-                Descriptor::Enum(descriptor) => {
-                    generate_enum(&resolver, type_path, descriptor, writer)?
-                }
+                Descriptor::Enum(descriptor) => generate_enum(
+                    &resolver,
+                    type_path,
+                    descriptor,
+                    writer,
+                    self.use_integers_for_enums,
+                )?,
                 Descriptor::Message(descriptor) => {
                     if let Some(message) = resolve_message(&self.descriptors, descriptor) {
                         generate_message(

pbjson-test/Cargo.toml

@@ -16,6 +16,7 @@ serde = { version = "1.0", features = ["derive"] }
 ignore-unknown-fields = []
 btree = []
 emit-fields = []
+use-integers-for-enums = []
 
 [dev-dependencies]
 chrono = "0.4"

pbjson-test/build.rs

@@ -54,6 +54,10 @@ fn main() -> Result<()> {
         builder.emit_fields();
     }
 
+    if cfg!(feature = "use-integers-for-enums") {
+        builder.use_integers_for_enums();
+    }
+
     builder.build(&[".test"])?;
 
     Ok(())

pbjson-test/src/lib.rs

@@ -46,6 +46,19 @@ mod tests {
     use pbjson_types::{Duration, Timestamp};
     use test::syntax3::*;
 
+    fn verify_encode(decoded: &KitchenSink, expected: &str) {
+        assert_eq!(serde_json::to_string(&decoded).unwrap().as_str(), expected);
+    }
+
+    fn verify_decode(decoded: &KitchenSink, expected: &str) {
+        assert_eq!(decoded, &serde_json::from_str(expected).unwrap());
+    }
+
+    fn verify(decoded: &KitchenSink, expected: &str) {
+        verify_encode(decoded, expected);
+        verify_decode(decoded, expected);
+    }
+
     #[test]
     #[cfg(not(feature = "ignore-unknown-fields"))]
     fn test_unknown_field_error() {
@@ -96,22 +109,18 @@ mod tests {
     }
 
     #[test]
-    #[cfg(not(feature = "emit-fields"))]
-    fn test_kitchen_sink() {
+    #[cfg(all(not(feature = "emit-fields"), feature = "use-integers-for-enums"))]
+    fn test_use_integers_for_enums() {
         let mut decoded: KitchenSink = serde_json::from_str("{}").unwrap();
+        assert_eq!(serde_json::to_string(&decoded).unwrap().as_str(), "{}");
+        decoded.value = kitchen_sink::Value::A as i32;
+        verify(&decoded, r#"{"value":45}"#);
+    }
 
-        let verify_encode = |decoded: &KitchenSink, expected: &str| {
-            assert_eq!(serde_json::to_string(&decoded).unwrap().as_str(), expected);
-        };
-
-        let verify_decode = |decoded: &KitchenSink, expected: &str| {
-            assert_eq!(decoded, &serde_json::from_str(expected).unwrap());
-        };
-
-        let verify = |decoded: &KitchenSink, expected: &str| {
-            verify_encode(decoded, expected);
-            verify_decode(decoded, expected);
-        };
+    #[test]
+    #[cfg(not(any(feature = "emit-fields", feature = "use-integers-for-enums")))]
+    fn test_kitchen_sink() {
+        let mut decoded: KitchenSink = serde_json::from_str("{}").unwrap();
 
         verify(&decoded, "{}");
         decoded.i32 = 24;