Merge pull request #556 from Shopify/el-fix-sorbet-signature-for-generic-jobs

Update sorbet compiler so that it correctly handles generic jobs

Author: emilylin0

lib/tapioca/dsl/compilers/job_iteration.rb

@@ -55,10 +55,18 @@ def decorate
             # Sorbet expects optional keyword arguments to be after required keyword arguments.
             expanded_parameters.sort_by! { |typed_param| PARAM_TYPES_IN_ORDER.index(typed_param.param.class) }
 
+            number_of_generic_type_members = Tapioca::Runtime::GenericTypeRegistry.lookup_type_variables(constant)&.size
+
+            returned_job_class = if number_of_generic_type_members&.nonzero?
+              "#{constant_name}[#{number_of_generic_type_members.times.map { "T.untyped" }.join(", ")}]"
+            else
+              constant_name
+            end
+
             job.create_method(
               "perform_later",
-              parameters: perform_later_parameters(expanded_parameters, constant_name),
-              return_type: "T.any(#{constant_name}, FalseClass)",
+              parameters: perform_later_parameters(expanded_parameters, returned_job_class),
+              return_type: "T.any(#{returned_job_class}, FalseClass)",
               class_method: true,
             )
 
@@ -83,15 +91,15 @@ def decorate
         sig do
           params(
             parameters: T::Array[RBI::TypedParam],
-            constant_name: T.nilable(String),
+            returned_job_class: String,
           ).returns(T::Array[RBI::TypedParam])
         end
-        def perform_later_parameters(parameters, constant_name)
+        def perform_later_parameters(parameters, returned_job_class)
           if ::Gem::Requirement.new(">= 7.0").satisfied_by?(::ActiveJob.gem_version)
             parameters.reject! { |typed_param| RBI::BlockParam === typed_param.param }
             parameters + [create_block_param(
               "block",
-              type: "T.nilable(T.proc.params(job: #{constant_name}).void)",
+              type: "T.nilable(T.proc.params(job: #{returned_job_class}).void)",
             )]
           else
             parameters

test/tapioca/dsl/compilers/job_iteration_test.rb

@@ -402,6 +402,87 @@ def perform_now(shop_ids:, profile_ids:, extension_ids:, foo:, bar:); end
           RBI
           assert_equal(expected, rbi_for(:NotifyJob))
         end
+
+        def test_generates_correct_rbi_file_for_generic_jobs_with_one_generic_variable
+          add_ruby_file("job.rb", <<~RUBY)
+            class GenericJob < ActiveJob::Base
+              include JobIteration::Iteration
+              extend T::Sig
+              extend T::Generic
+
+              Elem = type_member
+
+              sig { params(val: Elem, cursor: T.untyped).returns(T::Array[T.untyped]) }
+              def build_enumerator(val:, cursor:)
+                # ...
+              end
+            end
+          RUBY
+
+          expected = template(<<~RBI)
+            # typed: strong
+
+            class GenericJob
+              sig { params(val: Elem).void }
+              def perform(val:); end
+
+              class << self
+            <% if rails_version(">= 7.0") %>
+                sig { params(val: Elem, block: T.nilable(T.proc.params(job: GenericJob[T.untyped]).void)).returns(T.any(GenericJob[T.untyped], FalseClass)) }
+                def perform_later(val:, &block); end
+            <% else %>
+                sig { params(val: Elem).returns(T.any(GenericJob[T.untyped], FalseClass)) }
+                def perform_later(val:); end
+            <% end %>
+
+                sig { params(val: Elem).returns(T.any(NilClass, Exception)) }
+                def perform_now(val:); end
+              end
+            end
+          RBI
+          assert_equal(expected, rbi_for(:GenericJob))
+        end
+
+        def test_generates_correct_rbi_file_for_generic_jobs_with_multiple_generic_variables
+          add_ruby_file("job.rb", <<~RUBY)
+            class GenericJob < ActiveJob::Base
+              include JobIteration::Iteration
+              extend T::Sig
+              extend T::Generic
+
+              Elem = type_member
+              SomeValue = type_member
+
+              sig { params(val: Elem, some_value: SomeValue, cursor: T.untyped).returns(T::Array[T.untyped]) }
+              def build_enumerator(val:, some_value:, cursor:)
+                # ...
+              end
+            end
+          RUBY
+
+          expected = template(<<~RBI)
+            # typed: strong
+
+            class GenericJob
+              sig { params(val: Elem, some_value: SomeValue).void }
+              def perform(val:, some_value:); end
+
+              class << self
+            <% if rails_version(">= 7.0") %>
+                sig { params(val: Elem, some_value: SomeValue, block: T.nilable(T.proc.params(job: GenericJob[T.untyped, T.untyped]).void)).returns(T.any(GenericJob[T.untyped, T.untyped], FalseClass)) }
+                def perform_later(val:, some_value:, &block); end
+            <% else %>
+                sig { params(val: Elem, some_value: SomeValue).returns(T.any(GenericJob[T.untyped, T.untyped], FalseClass)) }
+                def perform_later(val:, some_value:); end
+            <% end %>
+
+                sig { params(val: Elem, some_value: SomeValue).returns(T.any(NilClass, Exception)) }
+                def perform_now(val:, some_value:); end
+              end
+            end
+          RBI
+          assert_equal(expected, rbi_for(:GenericJob))
+        end
       end
     end
   end