From 5bf728d4d048aa46cb5e9585660a0fbb8da0944d Mon Sep 17 00:00:00 2001
From: Andrew Novoselac <andrew.novoselac@shopify.com>
Date: Wed, 11 May 2022 19:27:55 -0400
Subject: [PATCH] Add support for :oneof fields on protobuf messages

---
 lib/tapioca/dsl/compilers/protobuf.rb       | 14 ++++++++++++
 spec/tapioca/dsl/compilers/protobuf_spec.rb | 24 +++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/lib/tapioca/dsl/compilers/protobuf.rb b/lib/tapioca/dsl/compilers/protobuf.rb
index 93e3eea35..da44b8221 100644
--- a/lib/tapioca/dsl/compilers/protobuf.rb
+++ b/lib/tapioca/dsl/compilers/protobuf.rb
@@ -83,6 +83,7 @@ def decorate
               create_type_members(klass, "Key", "Value")
             else
               descriptor = T.let(T.unsafe(constant).descriptor, Google::Protobuf::Descriptor)
+              descriptor.each_oneof { |oneof| create_oneof_method(klass, oneof) }
               fields = descriptor.map { |desc| create_descriptor_method(klass, desc) }
               fields.sort_by!(&:name)
 
@@ -216,6 +217,19 @@ def create_descriptor_method(klass, desc)
 
           field
         end
+
+        sig do
+          params(
+            klass: RBI::Scope,
+            desc: Google::Protobuf::OneofDescriptor
+          ).void
+        end
+        def create_oneof_method(klass, desc)
+          klass.create_method(
+            desc.name,
+            return_type: "T.nilable(Symbol)"
+          )
+        end
       end
     end
   end
diff --git a/spec/tapioca/dsl/compilers/protobuf_spec.rb b/spec/tapioca/dsl/compilers/protobuf_spec.rb
index 9996c3704..4436dcc89 100644
--- a/spec/tapioca/dsl/compilers/protobuf_spec.rb
+++ b/spec/tapioca/dsl/compilers/protobuf_spec.rb
@@ -411,6 +411,30 @@ def ShopName=(value); end
 
               assert_equal(expected, rbi_for(:Cart))
             end
+
+            it "generates methods in RBI files with oneof fields" do
+              add_ruby_file("protobuf.rb", <<~RUBY)
+                Google::Protobuf::DescriptorPool.generated_pool.build do
+                  add_file("cart.proto", :syntax => :proto3) do
+                    add_message "MyCart" do
+                      oneof :contact_info do
+                        optional :phone_number, :int32, 1
+                        optional :email, :string, 2
+                      end
+                    end
+                  end
+                end
+
+                Cart = Google::Protobuf::DescriptorPool.generated_pool.lookup("MyCart").msgclass
+              RUBY
+
+              rbi_output = rbi_for(:Cart)
+
+              assert_includes(rbi_output, indented(<<~RBI, 2))
+                sig { returns(T.nilable(Symbol)) }
+                def contact_info; end
+              RBI
+            end
           end
         end
       end