From da45252c8f274fb41014a32e7d26678b9006b19d Mon Sep 17 00:00:00 2001
From: Alex Tan <atan@triumphpay.com>
Date: Mon, 8 Apr 2024 19:16:42 +0700
Subject: [PATCH] Make id type non-nilable

---
 lib/tapioca/dsl/compilers/active_record_relations.rb     | 1 +
 lib/tapioca/helpers/rbi_helper.rb                        | 9 +++++++++
 .../dsl/compilers/active_record_relations_spec.rb        | 4 ++--
 spec/tapioca/helpers/rbi_helper_spec.rb                  | 9 +++++++++
 4 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/lib/tapioca/dsl/compilers/active_record_relations.rb b/lib/tapioca/dsl/compilers/active_record_relations.rb
index 3ff9ef9666..d0a070d282 100644
--- a/lib/tapioca/dsl/compilers/active_record_relations.rb
+++ b/lib/tapioca/dsl/compilers/active_record_relations.rb
@@ -676,6 +676,7 @@ def create_common_methods
               if constant.table_exists?
                 primary_key_type = constant.type_for_attribute(constant.primary_key)
                 type = Tapioca::Dsl::Helpers::ActiveModelTypeHelper.type_for(primary_key_type)
+                type = RBIHelper.as_non_nilable_type(type)
                 id_types << type if type != "T.untyped"
               end
 
diff --git a/lib/tapioca/helpers/rbi_helper.rb b/lib/tapioca/helpers/rbi_helper.rb
index 6ef444a148..a2249280a7 100644
--- a/lib/tapioca/helpers/rbi_helper.rb
+++ b/lib/tapioca/helpers/rbi_helper.rb
@@ -96,6 +96,15 @@ def as_nilable_type(type)
       end
     end
 
+    sig { params(type: String).returns(String) }
+    def as_non_nilable_type(type)
+      if type.match(/\A(?:::)?T.nilable\((.+)\)\z/)
+        T.must(::Regexp.last_match(1))
+      else
+        type
+      end
+    end
+
     sig { params(name: String).returns(T::Boolean) }
     def valid_method_name?(name)
       # try to parse a method definition with this name
diff --git a/spec/tapioca/dsl/compilers/active_record_relations_spec.rb b/spec/tapioca/dsl/compilers/active_record_relations_spec.rb
index 72cc8df0d4..0a9c28ad9a 100644
--- a/spec/tapioca/dsl/compilers/active_record_relations_spec.rb
+++ b/spec/tapioca/dsl/compilers/active_record_relations_spec.rb
@@ -130,8 +130,8 @@ def fifth; end
                     sig { returns(::Post) }
                     def fifth!; end
 
-                    sig { params(args: T.any(String, Symbol, ::ActiveSupport::Multibyte::Chars, T::Boolean, BigDecimal, Numeric, ::ActiveRecord::Type::Binary::Data, ::ActiveRecord::Type::Time::Value, Date, Time, ::ActiveSupport::Duration, T::Class[T.anything], T.nilable(::CustomId))).returns(::Post) }
-                    sig { params(args: T::Array[T.any(String, Symbol, ::ActiveSupport::Multibyte::Chars, T::Boolean, BigDecimal, Numeric, ::ActiveRecord::Type::Binary::Data, ::ActiveRecord::Type::Time::Value, Date, Time, ::ActiveSupport::Duration, T::Class[T.anything], T.nilable(::CustomId))]).returns(T::Enumerable[::Post]) }
+                    sig { params(args: T.any(String, Symbol, ::ActiveSupport::Multibyte::Chars, T::Boolean, BigDecimal, Numeric, ::ActiveRecord::Type::Binary::Data, ::ActiveRecord::Type::Time::Value, Date, Time, ::ActiveSupport::Duration, T::Class[T.anything], ::CustomId)).returns(::Post) }
+                    sig { params(args: T::Array[T.any(String, Symbol, ::ActiveSupport::Multibyte::Chars, T::Boolean, BigDecimal, Numeric, ::ActiveRecord::Type::Binary::Data, ::ActiveRecord::Type::Time::Value, Date, Time, ::ActiveSupport::Duration, T::Class[T.anything], ::CustomId)]).returns(T::Enumerable[::Post]) }
                     def find(args); end
 
                     sig { params(args: T.untyped).returns(T.nilable(::Post)) }
diff --git a/spec/tapioca/helpers/rbi_helper_spec.rb b/spec/tapioca/helpers/rbi_helper_spec.rb
index 50add629ea..3a9af1d313 100644
--- a/spec/tapioca/helpers/rbi_helper_spec.rb
+++ b/spec/tapioca/helpers/rbi_helper_spec.rb
@@ -7,6 +7,15 @@ class Tapioca::RBIHelperSpec < Minitest::Spec
   include Tapioca::RBIHelper
 
   describe Tapioca::RBIHelper do
+    specify "as_non_nilable_type removes T.nilable() and ::T.nilable() if it's the outermost part of the string" do
+      T.bind(self, T.untyped)
+
+      assert_equal(as_non_nilable_type("T.nilable(String)"), "String")
+      assert_equal(as_non_nilable_type("::T.nilable(String)"), "String")
+      assert_equal(as_non_nilable_type("String"), "String")
+      assert_equal(as_non_nilable_type("T.any(T.nilable(String), Integer)"), "T.any(T.nilable(String), Integer)")
+    end
+
     it "accepts valid method names" do
       [
         "f",