Merge pull request #1823 from bdewater/ar-relations-extract-associated

Fix return type of ActiveRecord::QueryMethods#extract_associated
Shopify · Mar 25, 2024 · 265bea2 · 265bea2
2 parents e3e5e6a + f4104b6
commit 265bea2
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 11 deletions.
diff --git a/lib/tapioca/dsl/compilers/active_record_relations.rb b/lib/tapioca/dsl/compilers/active_record_relations.rb
@@ -559,13 +559,29 @@ def create_relation_methods
           )
 
           QUERY_METHODS.each do |method_name|
-            create_relation_method(
-              method_name,
-              parameters: [
-                create_rest_param("args", type: "T.untyped"),
-                create_block_param("blk", type: "T.untyped"),
-              ],
-            )
+            case method_name
+            when :extract_associated
+              parameters = [create_param("association", type: "Symbol")]
+              return_type = "T::Array[T.untyped]"
+              relation_methods_module.create_method(
+                method_name.to_s,
+                parameters: parameters,
+                return_type: return_type,
+              )
+              association_relation_methods_module.create_method(
+                method_name.to_s,
+                parameters: parameters,
+                return_type: return_type,
+              )
+            else
+              create_relation_method(
+                method_name,
+                parameters: [
+                  create_rest_param("args", type: "T.untyped"),
+                  create_block_param("blk", type: "T.untyped"),
+                ],
+              )
+            end
           end
         end
 

diff --git a/spec/tapioca/dsl/compilers/active_record_relations_spec.rb b/spec/tapioca/dsl/compilers/active_record_relations_spec.rb
@@ -282,8 +282,8 @@ def excluding(*args, &blk); end
                     sig { params(args: T.untyped, blk: T.untyped).returns(PrivateAssociationRelation) }
                     def extending(*args, &blk); end
 
-                    sig { params(args: T.untyped, blk: T.untyped).returns(PrivateAssociationRelation) }
-                    def extract_associated(*args, &blk); end
+                    sig { params(association: Symbol).returns(T::Array[T.untyped]) }
+                    def extract_associated(association); end
 
                     sig { params(args: T.untyped, blk: T.untyped).returns(PrivateAssociationRelation) }
                     def from(*args, &blk); end
@@ -453,8 +453,8 @@ def excluding(*args, &blk); end
                     sig { params(args: T.untyped, blk: T.untyped).returns(PrivateRelation) }
                     def extending(*args, &blk); end
 
-                    sig { params(args: T.untyped, blk: T.untyped).returns(PrivateRelation) }
-                    def extract_associated(*args, &blk); end
+                    sig { params(association: Symbol).returns(T::Array[T.untyped]) }
+                    def extract_associated(association); end
 
                     sig { params(args: T.untyped, blk: T.untyped).returns(PrivateRelation) }
                     def from(*args, &blk); end