prost-build: generate setters for oneof fields

tokio-rs · Sep 1, 2023 · 1fb7926 · 1fb7926
1 parent 8e7a2cf
commit 1fb7926
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 3 deletions.
diff --git a/prost-build/src/code_generator.rs b/prost-build/src/code_generator.rs
@@ -673,7 +673,7 @@ impl<'a> CodeGenerator<'a> {
         message_name: &str,
         fq_message_name: &str,
         fields: &[(FieldDescriptorProto, usize)],
-        _oneof_fields: &[OneofField],
+        oneof_fields: &[OneofField],
     ) {
         debug!("  builder: {:?}", message_name);
 
@@ -708,7 +708,9 @@ impl<'a> CodeGenerator<'a> {
             self.append_builder_field_setter(fq_message_name, field);
         }
 
-        // TODO: generate oneof field setters
+        for oneof in oneof_fields {
+            self.append_builder_oneof_field_setter(&oneof.desc);
+        }
 
         self.push_indent();
         self.buf.push_str("#[inline]\n");
@@ -790,6 +792,35 @@ impl<'a> CodeGenerator<'a> {
         self.buf.push_str("}\n");
     }
 
+    fn append_builder_oneof_field_setter(&mut self, oneof: &OneofDescriptorProto) {
+        let field_name = to_snake(oneof.name());
+        let ty = to_upper_camel(oneof.name());
+
+        debug!("    oneof field setter: {:?}(arg: {:?})", field_name, ty);
+
+        self.push_indent();
+        self.buf.push_str("#[inline]\n");
+        self.push_indent();
+        self.buf.push_str("pub fn ");
+        self.buf.push_str(&field_name);
+        self.buf.push_str("(mut self, value: ");
+        self.buf.push_str(&ty);
+        self.buf.push_str(") -> Self {\n");
+        self.depth += 1;
+
+        self.push_indent();
+        self.buf.push_str("self.inner.");
+        self.buf.push_str(&field_name);
+        self.buf.push_str(" = Some(value);\n");
+
+        self.push_indent();
+        self.buf.push_str("self\n");
+
+        self.depth -= 1;
+        self.push_indent();
+        self.buf.push_str("}\n");
+    }
+
     fn location(&self) -> Option<&Location> {
         let source_info = self.source_info.as_ref()?;
         let idx = source_info

diff --git a/tests/src/builders.proto b/tests/src/builders.proto
@@ -21,4 +21,8 @@ message Zoo {
   Evolved message_field = 1;
   AnEnum enum_field = 2;
   EmptyForNow boxed_field = 3;
+  oneof oneof_field {
+    int32 a = 4;
+    string b = 5;
+  }
 }
diff --git a/tests/src/builders.rs b/tests/src/builders.rs
@@ -35,7 +35,7 @@ pub mod builders {
 
 #[cfg(test)]
 mod tests {
-    use super::builders::{AnEnum, EmptyForNow, Evolved, Zoo};
+    use super::builders::{zoo, AnEnum, EmptyForNow, Evolved, Zoo};
 
     use alloc::boxed::Box;
     use alloc::vec::Vec;
@@ -85,4 +85,10 @@ mod tests {
         let msg = Zoo::builder().boxed_field(EmptyForNow::default()).build();
         assert_eq!(msg.boxed_field, Some(Box::new(EmptyForNow::default())));
     }
+
+    #[test]
+    fn oneof_field_setter_takes_enum_values() {
+        let msg = Zoo::builder().oneof_field(zoo::OneofField::A(42)).build();
+        assert_eq!(msg.oneof_field, Some(zoo::OneofField::A(42)));
+    }
 }