Pass an option to ActiveRecordAssociations compiler for how to generate association types

lib/tapioca/dsl/compilers/active_record_associations.rb

@@ -13,19 +13,32 @@ module Compilers
       # This compiler is only responsible for defining the methods that would be created for the associations that
       # are defined in the Active Record model.
       #
+      # This compiler accepts a `ActiveRecordAssociationTypes` option that can be used to specify
+      # how the types of `belongs_to` and `has_one` associations should be generated. The option can be one of the
+      # following:
+      #  - `nilable (_default_)`: All association methods will be generated with `T.nilable` return types. This is
+      #  strictly the most correct way to type the methods, but it can make working with the models more cumbersome, as
+      #  you will have to handle the `nil` cases explicitly using `T.must` or the safe navigation operator `&.`, even
+      #  for valid persisted models.
+      #  - `persisted`: The methods will be generated with the type that matches validations on the association. If
+      #  there is a `required: true` or `optional: false`, then the types will be generated as non-nilable. This mode
+      #  basically treats each model as if it was a valid and persisted model. Note that this makes typing Active Record
+      #  models easier, but does not match the behaviour of non-persisted or invalid models, which can have `nil`
+      #  associations.
+      #
       # For example, with the following model class:
       #
       # ~~~rb
       # class Post < ActiveRecord::Base
       #   belongs_to :category
       #   has_many :comments
-      #   has_one :author, class_name: "User"
+      #   has_one :author, class_name: "User", optional: false
       #
       #   accepts_nested_attributes_for :category, :comments, :author
       # end
       # ~~~
       #
-      # this compiler will produce the following methods in the RBI file
+      # this compiler will produce, by default, the following methods in the RBI file
       # `post.rbi`:
       #
       # ~~~rbi
@@ -101,6 +114,16 @@ module Compilers
       #   end
       # end
       # ~~~
+      # If `ActiveRecordAssociationTypes` is `persisted`, the `author` method will be generated as:
+      # ~~~rbi
+      #     sig { returns(::User) }
+      #     def author; end
+      # ~~~
+      # and if the option is set to `untyped`, the `author` method will be generated as:
+      # ~~~rbi
+      #     sig { returns(T.untyped) }
+      #     def author; end
+      # ~~~
       class ActiveRecordAssociations < Compiler
         extend T::Sig
         include Helpers::ActiveRecordConstantsHelper
@@ -121,6 +144,50 @@ def initialize(class_name)
           end
         end
 
+        class AssociationTypeOption < T::Enum
+          extend T::Sig
+
+          enums do
+            Nilable = new("nilable")
+            Persisted = new("persisted")
+          end
+
+          class << self
+            extend T::Sig
+
+            sig do
+              params(
+                options: T::Hash[String, T.untyped],
+                block: T.proc.params(value: String, default_association_type_option: AssociationTypeOption).void,
+              ).returns(AssociationTypeOption)
+            end
+            def from_options(options, &block)
+              column_type_option = Nilable
+              value = options["ActiveRecordAssociationTypes"]
+
+              if value
+                if has_serialized?(value)
+                  column_type_option = from_serialized(value)
+                else
+                  block.call(value, column_type_option)
+                end
+              end
+
+              column_type_option
+            end
+          end
+
+          sig { returns(T::Boolean) }
+          def persisted?
+            self == AssociationTypeOption::Persisted
+          end
+
+          sig { returns(T::Boolean) }
+          def nilable?
+            self == AssociationTypeOption::Nilable
+          end
+        end
+
         ConstantType = type_member { { fixed: T.class_of(ActiveRecord::Base) } }
 
         sig { override.void }
@@ -148,6 +215,19 @@ def gather_constants
 
         private
 
+        sig { returns(AssociationTypeOption) }
+        def association_type_option
+          @association_type_option ||= T.let(
+            AssociationTypeOption.from_options(options) do |value, default_association_type_option|
+              add_error(<<~MSG.strip)
+                Unknown value for compiler option `ActiveRecordAssociationTypes` given: `#{value}`.
+                Proceeding with the default value: `#{default_association_type_option.serialize}`.
+              MSG
+            end,
+            T.nilable(AssociationTypeOption),
+          )
+        end
+
         sig { params(mod: RBI::Scope).void }
         def populate_nested_attribute_writers(mod)
           constant.nested_attributes_options.keys.each do |association_name|
@@ -187,7 +267,7 @@ def populate_associations(mod)
         end
         def populate_single_assoc_getter_setter(klass, association_name, reflection)
           association_class = type_for(reflection)
-          association_type = as_nilable_type(association_class)
+          association_type = single_association_type_for(reflection)
           association_methods_module = constant.generated_association_methods
 
           klass.create_method(
@@ -292,6 +372,45 @@ def type_for(reflection)
           T.must(qualified_name_of(reflection.klass))
         end
 
+        sig do
+          params(
+            reflection: ReflectionType,
+          ).returns(String)
+        end
+        def single_association_type_for(reflection)
+          association_class = type_for(reflection)
+          return as_nilable_type(association_class) unless association_type_option.persisted?
+
+          if has_one_and_required_reflection?(reflection) || belongs_to_and_non_optional_reflection?(reflection)
+            association_class
+          else
+            as_nilable_type(association_class)
+          end
+        end
+
+        # Note - one can do more here. If the association's attribute has an unconditional presence validation, it
+        # should also be considered required.
+        sig { params(reflection: ReflectionType).returns(T::Boolean) }
+        def has_one_and_required_reflection?(reflection)
+          return false unless reflection.has_one?
+          return false if reflection.options[:required].nil?
+
+          reflection.options[:required]
+        end
+
+        # Note - one can do more here. If the FK defining the belongs_to association is non-nullable at the DB level, or
+        # if the association's attribute has an unconditional presence validation, it should also be considered
+        # non-optional.
+        sig { params(reflection: ReflectionType).returns(T::Boolean) }
+        def belongs_to_and_non_optional_reflection?(reflection)
+          return false unless reflection.belongs_to?
+
+          required_by_default = !!reflection.active_record.belongs_to_required_by_default
+          return required_by_default if reflection.options[:optional].nil?
+
+          !reflection.options[:optional]
+        end
+
         sig do
           params(
             reflection: ReflectionType,

manual/compiler_activerecordassociations.md

@@ -5,19 +5,32 @@
 This compiler is only responsible for defining the methods that would be created for the associations that
 are defined in the Active Record model.
 
+This compiler accepts a `ActiveRecordAssociationTypes` option that can be used to specify
+how the types of `belongs_to` and `has_one` associations should be generated. The option can be one of the
+following:
+ - `nilable (_default_)`: All association methods will be generated with `T.nilable` return types. This is
+ strictly the most correct way to type the methods, but it can make working with the models more cumbersome, as
+ you will have to handle the `nil` cases explicitly using `T.must` or the safe navigation operator `&.`, even
+ for valid persisted models.
+ - `persisted`: The methods will be generated with the type that matches validations on the association. If
+ there is a `required: true` or `optional: false`, then the types will be generated as non-nilable. This mode
+ basically treats each model as if it was a valid and persisted model. Note that this makes typing Active Record
+ models easier, but does not match the behaviour of non-persisted or invalid models, which can have `nil`
+ associations.
+
 For example, with the following model class:
 
 ~~~rb
 class Post < ActiveRecord::Base
   belongs_to :category
   has_many :comments
-  has_one :author, class_name: "User"
+  has_one :author, class_name: "User", optional: false
 
   accepts_nested_attributes_for :category, :comments, :author
 end
 ~~~
 
-this compiler will produce the following methods in the RBI file
+this compiler will produce, by default, the following methods in the RBI file
 `post.rbi`:
 
 ~~~rbi
@@ -93,3 +106,13 @@ class Post
   end
 end
 ~~~
+If `ActiveRecordAssociationTypes` is `persisted`, the `author` method will be generated as:
+~~~rbi
+    sig { returns(::User) }
+    def author; end
+~~~
+and if the option is set to `untyped`, the `author` method will be generated as:
+~~~rbi
+    sig { returns(T.untyped) }
+    def author; end
+~~~