From 30c373e3e3db99e15977fe9a18d8ca65a460f706 Mon Sep 17 00:00:00 2001
From: Alexander Momchilov <alexander.momchilov@shopify.com>
Date: Fri, 14 Jun 2024 14:57:01 -0400
Subject: [PATCH] Support array types in `typed_store` compiler

---
 .../compilers/active_record_typed_store.rb    | 23 ++++++--
 .../active_record_typed_store_spec.rb         | 54 +++++++++++++++++++
 2 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/lib/tapioca/dsl/compilers/active_record_typed_store.rb b/lib/tapioca/dsl/compilers/active_record_typed_store.rb
index 59fd33922..91ded10b1 100644
--- a/lib/tapioca/dsl/compilers/active_record_typed_store.rb
+++ b/lib/tapioca/dsl/compilers/active_record_typed_store.rb
@@ -98,8 +98,7 @@ def decorate
             stores.values.each do |store_data|
               store_data.accessors.each do |accessor, name|
                 field = store_data.fields.fetch(accessor)
-                type = type_for(field.type_sym)
-                type = as_nilable_type(type) if field.null
+                type = type_for(field)
                 name ||= field.name # support < 1.5.0
 
                 generate_methods(store_accessors_module, name.to_s, type)
@@ -136,9 +135,23 @@ def gather_constants
           T::Hash[Symbol, String],
         )
 
-        sig { params(attr_type: Symbol).returns(String) }
-        def type_for(attr_type)
-          TYPES.fetch(attr_type, "T.untyped")
+        sig { params(field: ActiveRecord::TypedStore::Field).returns(String) }
+        def type_for(field)
+          type = TYPES.fetch(field.type_sym, "T.untyped")
+
+          type = if field.array
+            # `null: false` applies to the array itself, not the elements, which are always nilable.
+            # https://github.com/byroot/activerecord-typedstore/blob/2f3fb98/spec/support/models.rb#L46C34-L46C45
+            # https://github.com/byroot/activerecord-typedstore/blob/2f3fb98/spec/active_record/typed_store_spec.rb#L854-L857
+            nilable_element_type = as_nilable_type(type)
+            "T::Array[#{nilable_element_type}]"
+          else
+            type
+          end
+
+          type = as_nilable_type(type) if field.null
+
+          type
         end
 
         sig do
diff --git a/spec/tapioca/dsl/compilers/active_record_typed_store_spec.rb b/spec/tapioca/dsl/compilers/active_record_typed_store_spec.rb
index b32444d84..a524e5a28 100644
--- a/spec/tapioca/dsl/compilers/active_record_typed_store_spec.rb
+++ b/spec/tapioca/dsl/compilers/active_record_typed_store_spec.rb
@@ -428,6 +428,60 @@ def rate=(rate); end
               assert_includes(rbi_for(:Post), expected)
             end
 
+            it "generates methods with nilable Array type of nilable elements for attributes marked as array: true" do
+              add_ruby_file("post.rb", <<~RUBY)
+                class Post < ActiveRecord::Base
+                  typed_store :metadata do |s|
+                    s.string(:comments, array: true, null: true)
+                  end
+                end
+              RUBY
+
+              expected = <<~RBI
+                # typed: strong
+
+                class Post
+                  include StoreAccessors
+
+                  module StoreAccessors
+                    sig { returns(T.nilable(T::Array[T.nilable(String)])) }
+                    def comments; end
+
+                    sig { params(comments: T.nilable(T::Array[T.nilable(String)])).returns(T.nilable(T::Array[T.nilable(String)])) }
+                    def comments=(comments); end
+              RBI
+
+              rbi = rbi_for(:Post)
+              assert_includes(rbi, expected)
+            end
+
+            it "generates methods with non-nilable Array type for attributes marked as array: true and null: false" do
+              add_ruby_file("post.rb", <<~RUBY)
+                class Post < ActiveRecord::Base
+                  typed_store :metadata do |s|
+                    s.string(:comments, array: true, null: false)
+                  end
+                end
+              RUBY
+
+              expected = <<~RBI
+                # typed: strong
+
+                class Post
+                  include StoreAccessors
+
+                  module StoreAccessors
+                    sig { returns(T::Array[T.nilable(String)]) }
+                    def comments; end
+
+                    sig { params(comments: T::Array[T.nilable(String)]).returns(T::Array[T.nilable(String)]) }
+                    def comments=(comments); end
+              RBI
+
+              rbi = rbi_for(:Post)
+              assert_includes(rbi, expected)
+            end
+
             it "generates methods with prefix and suffix" do
               add_ruby_file("post.rb", <<~RUBY)
                 class Post < ActiveRecord::Base