Fix logic for gather_constants and add test

Co-authored-by: Ufuk Kayserilioglu <[email protected]>
Shopify · Feb 23, 2023 · 620ce9e · 620ce9e
1 parent f887407
commit 620ce9e
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 6 deletions.
diff --git a/lib/tapioca/dsl/compilers/protobuf.rb b/lib/tapioca/dsl/compilers/protobuf.rb
@@ -162,13 +162,17 @@ def gather_constants
               T.cast(desc, Google::Protobuf::EnumDescriptor).enummodule
             end
 
-            abstract_message_const = ::Google::Protobuf.const_get(:AbstractMessage)
-            abstract_message_descendants = descendants_of(abstract_message_const)
+            results = if Google::Protobuf.const_defined?(:AbstractMessage)
+              abstract_message_const = ::Google::Protobuf.const_get(:AbstractMessage)
+              descendants_of(abstract_message_const) - [abstract_message_const]
+            else
+              T.cast(
+                ObjectSpace.each_object(marker).to_a,
+                T::Array[Module],
+              )
+            end
 
-            results = T.cast(
-              ObjectSpace.each_object(marker).to_a,
-              T::Array[Module],
-            ).concat(enum_modules).concat(abstract_message_descendants) - [abstract_message_const]
+            results = results.concat(enum_modules)
             results.any? ? results + [Google::Protobuf::RepeatedField, Google::Protobuf::Map] : []
           end
         end

diff --git a/spec/tapioca/dsl/compilers/protobuf_spec.rb b/spec/tapioca/dsl/compilers/protobuf_spec.rb
@@ -36,6 +36,24 @@ class ProtobufSpec < ::DslSpec
               assert_includes(gathered_constants, "Google::Protobuf::Map")
               assert_includes(gathered_constants, "Google::Protobuf::RepeatedField")
             end
+
+            it "skips AbstractMessage" do
+              add_ruby_file("content.rb", <<~RUBY)
+                Google::Protobuf::DescriptorPool.generated_pool.build do
+                  add_file("cart.proto", :syntax => :proto3) do
+                    add_message "MyCart" do
+                      optional :shop_id, :int32, 1
+                      optional :customer_id, :int32, 2
+                    end
+                  end
+                end
+
+                Cart = Google::Protobuf::DescriptorPool.generated_pool.lookup("MyCart").msgclass
+              RUBY
+
+              assert_equal(["Cart"], gathered_constants.reject { |constant| constant.start_with?("Google::Protobuf") })
+              refute_includes(gathered_constants, "Google::Protobuf::AbstractMessage")
+            end
           end
 
           describe "decorate" do