Merge pull request #1830 from stathis-alexander/ajs/correct-sigs-for-…

…active-record-sum generate overloaded sigs for active record sum
Shopify · Apr 19, 2024 · bfc8731 · bfc8731
2 parents e13dd7c + 93cb91e
commit bfc8731
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 16 deletions.
diff --git a/lib/tapioca/dsl/compilers/active_record_relations.rb b/lib/tapioca/dsl/compilers/active_record_relations.rb
@@ -822,13 +822,28 @@ def create_common_methods
                 return_type: "T.untyped",
               )
             when :sum
-              create_common_method(
-                "sum",
+              sigs = [
+                common_relation_methods_module.create_sig(
+                  parameters: { initial_value_or_column: "T.untyped" },
+                  return_type: "Numeric",
+                ),
+                common_relation_methods_module.create_sig(
+                  type_parameters: ["U"],
+                  parameters:
+                  {
+                    initial_value_or_column: "T.nilable(T.type_parameter(:U))",
+                    block: "T.proc.params(object: #{constant_name}).returns(T.type_parameter(:U))",
+                  },
+                  return_type: "T.type_parameter(:U)",
+                ),
+              ]
+              common_relation_methods_module.create_method_with_sigs(
+                method_name.to_s,
+                sigs: sigs,
                 parameters: [
-                  create_opt_param("column_name", type: "T.nilable(T.any(String, Symbol))", default: "nil"),
-                  create_block_param("block", type: "T.nilable(T.proc.params(record: T.untyped).returns(T.untyped))"),
+                  RBI::OptParam.new("initial_value_or_column", "nil"),
+                  RBI::BlockParam.new("block"),
                 ],
-                return_type: "Numeric",
               )
             end
           end

diff --git a/lib/tapioca/rbi_ext/model.rb b/lib/tapioca/rbi_ext/model.rb
@@ -128,14 +128,15 @@ def create_method_with_sigs(name, sigs:, parameters: [], class_method: false, vi
     sig do
       params(
         parameters: T::Hash[T.any(String, Symbol), String],
+        type_parameters: T::Array[String],
         return_type: String,
       ).returns(RBI::Sig)
     end
-    def create_sig(parameters:, return_type: "T.untyped")
+    def create_sig(parameters:, type_parameters: [], return_type: "T.untyped")
       params = parameters.map do |name, type|
         RBI::SigParam.new(name.to_s, type)
       end
-      RBI::Sig.new(params: params, return_type: return_type)
+      RBI::Sig.new(type_params: type_parameters, params: params, return_type: return_type)
     end
 
     private

diff --git a/sorbet/rbi/gems/rbi@0.1.10.rbi b/sorbet/rbi/gems/rbi@0.1.10.rbi
diff --git a/spec/tapioca/dsl/compilers/active_record_relations_spec.rb b/spec/tapioca/dsl/compilers/active_record_relations_spec.rb
@@ -231,8 +231,9 @@ def second_to_last!; end
                     def sole; end
 
                 <% end %>
-                    sig { params(column_name: T.nilable(T.any(String, Symbol)), block: T.nilable(T.proc.params(record: T.untyped).returns(T.untyped))).returns(Numeric) }
-                    def sum(column_name = nil, &block); end
+                    sig { params(initial_value_or_column: T.untyped).returns(Numeric) }
+                    sig { type_parameters(:U).params(initial_value_or_column: T.nilable(T.type_parameter(:U)), block: T.proc.params(object: ::Post).returns(T.type_parameter(:U))).returns(T.type_parameter(:U)) }
+                    def sum(initial_value_or_column = nil, &block); end
 
                     sig { params(limit: NilClass).returns(T.nilable(::Post)) }
                     sig { params(limit: Integer).returns(T::Array[::Post]) }
@@ -927,8 +928,9 @@ def second_to_last!; end
                     def sole; end
 
                 <% end %>
-                    sig { params(column_name: T.nilable(T.any(String, Symbol)), block: T.nilable(T.proc.params(record: T.untyped).returns(T.untyped))).returns(Numeric) }
-                    def sum(column_name = nil, &block); end
+                    sig { params(initial_value_or_column: T.untyped).returns(Numeric) }
+                    sig { type_parameters(:U).params(initial_value_or_column: T.nilable(T.type_parameter(:U)), block: T.proc.params(object: ::Post).returns(T.type_parameter(:U))).returns(T.type_parameter(:U)) }
+                    def sum(initial_value_or_column = nil, &block); end
 
                     sig { params(limit: NilClass).returns(T.nilable(::Post)) }
                     sig { params(limit: Integer).returns(T::Array[::Post]) }