Add types to generated RBI files for FrozenRecord

lib/tapioca/dsl/compilers/frozen_record.rb

@@ -77,8 +77,18 @@ def decorate
 
             record.create_module(module_name) do |mod|
               attributes.each do |attribute|
+                return_type = "T.untyped"
+                if constant.respond_to?(:attribute_types)
+                  attribute_type = T.let(
+                    T.unsafe(constant).attribute_types[attribute],
+                    ActiveModel::Type::Value
+                  )
+                  has_default = T.let(constant.default_attributes.key?(attribute), T::Boolean)
+                  return_type = type_for(attribute_type, has_default)
+                end
+
                 mod.create_method("#{attribute}?", return_type: "T::Boolean")
-                mod.create_method(attribute.to_s, return_type: "T.untyped")
+                mod.create_method(attribute.to_s, return_type: return_type)
               end
             end
 
@@ -95,6 +105,41 @@ def self.gather_constants
 
         private
 
+        sig { params(attribute_type_value: ::ActiveModel::Type::Value, has_default: T::Boolean).returns(::String) }
+        def type_for(attribute_type_value, has_default)
+          type = case attribute_type_value
+          when ActiveModel::Type::Boolean
+            "T::Boolean"
+          when ActiveModel::Type::Date
+            "::Date"
+          when ActiveModel::Type::DateTime, ActiveModel::Type::Time
+            "::DateTime"
+          when ActiveModel::Type::Decimal
+            "::BigDecimal"
+          when ActiveModel::Type::Float
+            "::Float"
+          when ActiveModel::Type::Integer
+            "::Integer"
+          when ActiveModel::Type::String
+            "::String"
+          else
+            other_type = attribute_type_value.type
+            case other_type
+            when :array
+              "::Array"
+            when :hash
+              "::Hash"
+            when :symbol
+              "::Symbol"
+            else
+              # we don't want untyped to be wrapped by T.nilable, so just return early
+              return "T.untyped"
+            end
+          end
+
+          has_default ? type : as_nilable_type(type)
+        end
+
         sig { params(record: RBI::Scope).void }
         def decorate_scopes(record)
           scopes = T.unsafe(constant).__tapioca_scope_names

spec/tapioca/dsl/compilers/frozen_record_spec.rb

@@ -97,6 +97,195 @@ def last_name?; end
               assert_equal(expected, rbi_for(:Student))
             end
 
+            it "can handle annotated fields" do
+              add_ruby_file("student.rb", <<~RUBY)
+                # typed: strong
+
+                class ArrayOfType < ActiveModel::Type::Value
+                  attr_reader :element_type
+
+                  def initialize(element_type:)
+                    super()
+                    @element_type = element_type
+                  end
+
+                  def type
+                    :array
+                  end
+                end
+
+                class HashOfType < ActiveModel::Type::Value
+                  attr_reader :key_type
+                  attr_reader :value_type
+
+                  def initialize(key_type:, value_type:)
+                    super()
+                    @key_type = key_type
+                    @value_type = value_type
+                  end
+
+                  def type
+                    :hash
+                  end
+                end
+
+                class SymbolType < ActiveModel::Type::Value
+                  def type
+                    :symbol
+                  end
+                end
+
+                ActiveModel::Type.register(:array_of_type, ArrayOfType)
+                ActiveModel::Type.register(:hash_of_type, HashOfType)
+                ActiveModel::Type.register(:symbol, SymbolType)
+
+                class Student < FrozenRecord::Base
+                  extend T::Sig
+                  include ActiveModel::Attributes
+
+                  # specifically missing the id field, should be untyped
+                  attribute :first_name, :string
+                  attribute :last_name, :string
+                  attribute :age, :integer
+                  attribute :location, :string
+                  attribute :is_cool_person, :boolean
+                  attribute :birth_date, :date
+                  attribute :updated_at, :time
+                  # custom attribute types
+                  attribute :favourite_foods, :array_of_type, element_type: :string
+                  attribute :skills, :hash_of_type, key_type: :symbol, value_type: :string
+                  # attribute with a default, shouldn't be nilable
+                  attribute :shirt_size, :symbol
+
+                  self.base_path = __dir__
+                  self.default_attributes = { shirt_size: :large }
+
+                  # Explicit method, shouldn't be in the RBI output
+                  sig { params(grain: Symbol).returns(String) }
+                  def area(grain:)
+                    parts = location.split(',').map(&:strip)
+                    case grain
+                    when :city
+                      parts[0]
+                    when :province
+                      parts[1]
+                    when :country
+                      parts[2]
+                    else
+                      location
+                    end
+                  end
+                end
+              RUBY
+
+              add_content_file("students.yml", <<~YAML)
+                - id: 1
+                  first_name: John
+                  last_name: Smith
+                  age: 19
+                  location: Ottawa, Ontario, Canada
+                  is_cool_person: no
+                  birth_date: 1867-07-01
+                  updated_at: 2014-02-24T19:08:06-05:00
+                  favourite_foods:
+                    - Pizza
+                  skills:
+                    backend: Ruby
+                    frontend: HTML
+                - id: 2
+                  first_name: Dan
+                  last_name:  Lord
+                  age: 20
+                  location: Toronto, Ontario, Canada
+                  is_cool_person: yes
+                  birth_date: 1967-07-01
+                  updated_at: 2015-02-24T19:08:06-05:00
+                  favourite_foods:
+                    - Tacos
+                  skills:
+                    backend: Ruby
+                    frontend: CSS
+              YAML
+
+              expected = <<~RBI
+                # typed: strong
+
+                class Student
+                  include FrozenRecordAttributeMethods
+
+                  module FrozenRecordAttributeMethods
+                    sig { returns(T.nilable(::Integer)) }
+                    def age; end
+
+                    sig { returns(T::Boolean) }
+                    def age?; end
+
+                    sig { returns(T.nilable(::Date)) }
+                    def birth_date; end
+
+                    sig { returns(T::Boolean) }
+                    def birth_date?; end
+
+                    sig { returns(T.nilable(::Array)) }
+                    def favourite_foods; end
+
+                    sig { returns(T::Boolean) }
+                    def favourite_foods?; end
+
+                    sig { returns(T.nilable(::String)) }
+                    def first_name; end
+
+                    sig { returns(T::Boolean) }
+                    def first_name?; end
+
+                    sig { returns(T.untyped) }
+                    def id; end
+
+                    sig { returns(T::Boolean) }
+                    def id?; end
+
+                    sig { returns(T.nilable(T::Boolean)) }
+                    def is_cool_person; end
+
+                    sig { returns(T::Boolean) }
+                    def is_cool_person?; end
+
+                    sig { returns(T.nilable(::String)) }
+                    def last_name; end
+
+                    sig { returns(T::Boolean) }
+                    def last_name?; end
+
+                    sig { returns(T.nilable(::String)) }
+                    def location; end
+
+                    sig { returns(T::Boolean) }
+                    def location?; end
+
+                    sig { returns(::Symbol) }
+                    def shirt_size; end
+
+                    sig { returns(T::Boolean) }
+                    def shirt_size?; end
+
+                    sig { returns(T.nilable(::Hash)) }
+                    def skills; end
+
+                    sig { returns(T::Boolean) }
+                    def skills?; end
+
+                    sig { returns(T.nilable(::DateTime)) }
+                    def updated_at; end
+
+                    sig { returns(T::Boolean) }
+                    def updated_at?; end
+                  end
+                end
+              RBI
+
+              assert_equal(expected, rbi_for(:Student))
+            end
+
             it "can handle frozen record scopes" do
               add_ruby_file("student.rb", <<~RUBY)
                 class Student < FrozenRecord::Base