Merge pull request #1860 from bdewater/ar-count-block

Update Active Record count sig to allow a block
Shopify · Apr 10, 2024 · 032e04f · 032e04f
2 parents 218b9c0 + 5bba1b3
commit 032e04f
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git a/lib/tapioca/dsl/compilers/active_record_relations.rb b/lib/tapioca/dsl/compilers/active_record_relations.rb
@@ -793,12 +793,23 @@ def create_common_methods
                 return_type: "Numeric",
               )
             when :count
-              create_common_method(
+              sigs = [
+                common_relation_methods_module.create_sig(
+                  parameters: { column_name: "T.nilable(T.any(String, Symbol))" },
+                  return_type: "Integer",
+                ),
+                common_relation_methods_module.create_sig(
+                  parameters: { column_name: "NilClass", block: "T.proc.params(object: #{constant_name}).void" },
+                  return_type: "Integer",
+                ),
+              ]
+              common_relation_methods_module.create_method_with_sigs(
                 "count",
+                sigs: sigs,
                 parameters: [
-                  create_opt_param("column_name", type: "T.untyped", default: "nil"),
+                  RBI::OptParam.new("column_name", "nil"),
+                  RBI::BlockParam.new("block"),
                 ],
-                return_type: "Integer",
               )
             when :ids
               create_common_method("ids", return_type: "Array")

diff --git a/spec/tapioca/dsl/compilers/active_record_relations_spec.rb b/spec/tapioca/dsl/compilers/active_record_relations_spec.rb
@@ -78,8 +78,9 @@ def build(attributes = nil, &block); end
                     sig { params(operation: Symbol, column_name: T.any(String, Symbol)).returns(Numeric) }
                     def calculate(operation, column_name); end
 
-                    sig { params(column_name: T.untyped).returns(Integer) }
-                    def count(column_name = nil); end
+                    sig { params(column_name: T.nilable(T.any(String, Symbol))).returns(Integer) }
+                    sig { params(column_name: NilClass, block: T.proc.params(object: ::Post).void).returns(Integer) }
+                    def count(column_name = nil, &block); end
 
                     sig { params(attributes: T.untyped, block: T.nilable(T.proc.params(object: ::Post).void)).returns(::Post) }
                     def create(attributes = nil, &block); end
@@ -769,8 +770,9 @@ def build(attributes = nil, &block); end
                     sig { params(operation: Symbol, column_name: T.any(String, Symbol)).returns(Numeric) }
                     def calculate(operation, column_name); end
 
-                    sig { params(column_name: T.untyped).returns(Integer) }
-                    def count(column_name = nil); end
+                    sig { params(column_name: T.nilable(T.any(String, Symbol))).returns(Integer) }
+                    sig { params(column_name: NilClass, block: T.proc.params(object: ::Post).void).returns(Integer) }
+                    def count(column_name = nil, &block); end
 
                     sig { params(attributes: T.untyped, block: T.nilable(T.proc.params(object: ::Post).void)).returns(::Post) }
                     def create(attributes = nil, &block); end