Add support for alias_attribute

chanzuckerberg · Sep 24, 2022 · 746e8c7 · 746e8c7
1 parent 447ecbe
commit 746e8c7
Show file tree

Hide file tree

Showing 15 changed files with 177 additions and 14 deletions.
diff --git a/lib/sorbet-rails/model_plugins/active_record_attribute.rb b/lib/sorbet-rails/model_plugins/active_record_attribute.rb
@@ -12,28 +12,33 @@ def generate(root)
 
     model_class_rbi = root.create_class(self.model_class_name)
     model_class_rbi.create_include(attribute_module_name)
-    model_defined_enums = @model_class.defined_enums
 
-    columns_hash.sort.each do |column_name, column_def|
-      if model_defined_enums.has_key?(column_name)
+    aliases_hash = model_class_columns_to_aliases
+    attributes_and_aliases_hash = model_class_attributes_and_aliases(columns_hash, aliases_hash)
+    enum_attributes_and_aliases_hash = model_class_enum_attributes_and_aliases(aliases_hash)
+
+    attributes_and_aliases_hash.each do |attribute_name, column_def|
+      if enum_attributes_and_aliases_hash.has_key?(attribute_name)
         generate_enum_methods(
           root,
           model_class_rbi,
           attribute_module_rbi,
-          model_defined_enums,
-          column_name,
+          enum_attributes_and_aliases_hash,
+          attribute_name,
           column_def,
         )
-      elsif serialization_coder_for_column(column_name)
+      elsif serialization_coder_for_column(attribute_name)
         next # handled by the ActiveRecordSerializedAttribute plugin
       else
         column_type = type_for_column_def(column_def)
+
         attribute_module_rbi.create_method(
-          column_name.to_s,
+          attribute_name,
           return_type: column_type.to_s,
         )
+
         attribute_module_rbi.create_method(
-          "#{column_name}=",
+          "#{attribute_name}=",
           parameters: [
             Parameter.new("value", type: value_type_for_attr_writer(column_type))
           ],
@@ -42,12 +47,14 @@ def generate(root)
       end
 
       attribute_module_rbi.create_method(
-        "#{column_name}?",
+        "#{attribute_name}?",
         return_type: "T::Boolean",
       )
     end
   end
 
+  private
+
   sig {
     params(
       root: Parlour::RbiGenerator::Namespace,
@@ -144,4 +151,43 @@ def value_type_for_attr_writer(column_type)
     end
     ColumnType.new(base_type: type, nilable: column_type.nilable, array_type: column_type.array_type).to_s
   end
+
+  sig { returns(T::Hash[String, T::Array[String]]) }
+  def model_class_columns_to_aliases
+    T.unsafe(@model_class)
+      .attribute_aliases
+      .each_with_object({}) do |(alias_, column), hash|
+        hash[column] ||= []
+        hash[column] << alias_
+      end
+  end
+
+  sig {
+    params(
+      columns_hash: T::Hash[String, T.untyped],
+      aliases_hash: T::Hash[String, T::Array[String]],
+    ).returns(T::Hash[String, T.untyped])
+  }
+  def model_class_attributes_and_aliases(columns_hash, aliases_hash)
+    columns_hash
+      .sort
+      .each_with_object({}) do |(column_name, column_def), hash|
+        attributes = [column_name].union(aliases_hash[column_name] || [])
+        attributes.each do |attribute_name|
+          hash[attribute_name] = column_def
+        end
+      end
+  end
+
+  sig { params(aliases_hash: T::Hash[String, T::Array[String]]).returns(T::Hash[String, T.untyped]) }
+  def model_class_enum_attributes_and_aliases(aliases_hash)
+    @model_class
+      .defined_enums
+      .each_with_object({}) do |(column_name, column_def), hash|
+        attributes = [column_name].union(aliases_hash[column_name] || [])
+        attributes.each do |attribute_name|
+          hash[attribute_name] = column_def
+        end
+      end
+  end
 end
diff --git a/spec/generators/rails-template.rb b/spec/generators/rails-template.rb
@@ -92,6 +92,9 @@ class SpellBook < ApplicationRecord
         dark_art: 999,
       }
 
+      alias_attribute :title, :name
+      alias_attribute :book_category, :book_type
+
       scope :recent, -> { where('created_at > ?', 1.month.ago) }
     end
   RUBY

diff --git a/spec/generators/sorbet_test_cases.rb b/spec/generators/sorbet_test_cases.rb
@@ -8,6 +8,13 @@
 # -- model columns
 T.assert_type!(wizard.name, T.nilable(String))
 
+spell_book = wizard.spell_books.first!
+T.assert_type!(spell_book, SpellBook)
+
+# -- alias attributes
+T.assert_type!(spell_book.title, String)
+T.assert_type!(spell_book.book_category, String)
+
 # -- time/date columns
 T.assert_type!(wizard.created_at, ActiveSupport::TimeWithZone)
 T.assert_type!(wand.broken_at, T.nilable(Time))
@@ -157,7 +164,6 @@
 T.assert_type!(Wizard.all.empty?, T::Boolean)
 
 # Finder methods -- CollectionProxy
-spell_book = wizard.spell_books.first!
 spell_books = wizard.spell_books
 T.assert_type!(spell_books.exists?(name: 'Fantastic Beasts'), T::Boolean)
 T.assert_type!(spell_books.find(spell_book.id), SpellBook)

diff --git a/spec/support/v5.2/app/models/spell_book.rb b/spec/support/v5.2/app/models/spell_book.rb
@@ -14,5 +14,8 @@ class SpellBook < ApplicationRecord
     dark_art: 999,
   }
 
+  alias_attribute :title, :name
+  alias_attribute :book_category, :book_type
+
   scope :recent, -> { where('created_at > ?', 1.month.ago) }
 end
diff --git a/spec/support/v5.2/sorbet_test_cases.rb b/spec/support/v5.2/sorbet_test_cases.rb
@@ -8,6 +8,13 @@
 # -- model columns
 T.assert_type!(wizard.name, T.nilable(String))
 
+spell_book = wizard.spell_books.first!
+T.assert_type!(spell_book, SpellBook)
+
+# -- alias attributes
+T.assert_type!(spell_book.title, String)
+T.assert_type!(spell_book.book_category, String)
+
 # -- time/date columns
 T.assert_type!(wizard.created_at, ActiveSupport::TimeWithZone)
 T.assert_type!(wand.broken_at, T.nilable(Time))
@@ -161,7 +168,6 @@
 T.assert_type!(Wizard.all.empty?, T::Boolean)
 
 # Finder methods -- CollectionProxy
-spell_book = wizard.spell_books.first!
 spell_books = wizard.spell_books
 T.assert_type!(spell_books.exists?(name: 'Fantastic Beasts'), T::Boolean)
 T.assert_type!(spell_books.find(spell_book.id), SpellBook)

diff --git a/spec/support/v6.0/app/models/spell_book.rb b/spec/support/v6.0/app/models/spell_book.rb
@@ -14,5 +14,8 @@ class SpellBook < ApplicationRecord
     dark_art: 999,
   }
 
+  alias_attribute :title, :name
+  alias_attribute :book_category, :book_type
+
   scope :recent, -> { where('created_at > ?', 1.month.ago) }
 end
diff --git a/spec/support/v6.0/sorbet_test_cases.rb b/spec/support/v6.0/sorbet_test_cases.rb
@@ -8,6 +8,13 @@
 # -- model columns
 T.assert_type!(wizard.name, T.nilable(String))
 
+spell_book = wizard.spell_books.first!
+T.assert_type!(spell_book, SpellBook)
+
+# -- alias attributes
+T.assert_type!(spell_book.title, String)
+T.assert_type!(spell_book.book_category, String)
+
 # -- time/date columns
 T.assert_type!(wizard.created_at, ActiveSupport::TimeWithZone)
 T.assert_type!(wand.broken_at, T.nilable(Time))
@@ -161,7 +168,6 @@
 T.assert_type!(Wizard.all.empty?, T::Boolean)
 
 # Finder methods -- CollectionProxy
-spell_book = wizard.spell_books.first!
 spell_books = wizard.spell_books
 T.assert_type!(spell_books.exists?(name: 'Fantastic Beasts'), T::Boolean)
 T.assert_type!(spell_books.find(spell_book.id), SpellBook)

diff --git a/spec/support/v6.1/app/models/spell_book.rb b/spec/support/v6.1/app/models/spell_book.rb
@@ -14,5 +14,8 @@ class SpellBook < ApplicationRecord
     dark_art: 999,
   }
 
+  alias_attribute :title, :name
+  alias_attribute :book_category, :book_type
+
   scope :recent, -> { where('created_at > ?', 1.month.ago) }
 end
diff --git a/spec/support/v6.1/sorbet_test_cases.rb b/spec/support/v6.1/sorbet_test_cases.rb
@@ -8,6 +8,13 @@
 # -- model columns
 T.assert_type!(wizard.name, T.nilable(String))
 
+spell_book = wizard.spell_books.first!
+T.assert_type!(spell_book, SpellBook)
+
+# -- alias attributes
+T.assert_type!(spell_book.title, String)
+T.assert_type!(spell_book.book_category, String)
+
 # -- time/date columns
 T.assert_type!(wizard.created_at, ActiveSupport::TimeWithZone)
 T.assert_type!(wand.broken_at, T.nilable(Time))
@@ -163,7 +170,6 @@
 T.assert_type!(Wizard.all.empty?, T::Boolean)
 
 # Finder methods -- CollectionProxy
-spell_book = wizard.spell_books.first!
 spell_books = wizard.spell_books
 T.assert_type!(spell_books.exists?(name: 'Fantastic Beasts'), T::Boolean)
 T.assert_type!(spell_books.find(spell_book.id), SpellBook)

diff --git a/spec/support/v7.0/app/models/spell_book.rb b/spec/support/v7.0/app/models/spell_book.rb
@@ -14,5 +14,8 @@ class SpellBook < ApplicationRecord
     dark_art: 999,
   }
 
+  alias_attribute :title, :name
+  alias_attribute :book_category, :book_type
+
   scope :recent, -> { where('created_at > ?', 1.month.ago) }
 end
diff --git a/spec/support/v7.0/sorbet_test_cases.rb b/spec/support/v7.0/sorbet_test_cases.rb
@@ -8,6 +8,13 @@
 # -- model columns
 T.assert_type!(wizard.name, T.nilable(String))
 
+spell_book = wizard.spell_books.first!
+T.assert_type!(spell_book, SpellBook)
+
+# -- alias attributes
+T.assert_type!(spell_book.title, String)
+T.assert_type!(spell_book.book_category, String)
+
 # -- time/date columns
 T.assert_type!(wizard.created_at, ActiveSupport::TimeWithZone)
 T.assert_type!(wand.broken_at, T.nilable(Time))
@@ -162,7 +169,6 @@
 T.assert_type!(Wizard.all.empty?, T::Boolean)
 
 # Finder methods -- CollectionProxy
-spell_book = wizard.spell_books.first!
 spell_books = wizard.spell_books
 T.assert_type!(spell_books.exists?(name: 'Fantastic Beasts'), T::Boolean)
 T.assert_type!(spell_books.find(spell_book.id), SpellBook)

diff --git a/spec/test_data/v5.2/expected_spell_book.rbi b/spec/test_data/v5.2/expected_spell_book.rbi
@@ -37,6 +37,15 @@ module SpellBook::GeneratedAttributeMethods
   sig { returns(T::Boolean) }
   def book_type?; end
 
+  sig { returns(String) }
+  def book_category; end
+
+  sig { params(value: T.any(Integer, String, Symbol)).void }
+  def book_category=(value); end
+
+  sig { returns(T::Boolean) }
+  def book_category?; end
+
   sig { returns(Integer) }
   def id; end
 
@@ -55,6 +64,15 @@ module SpellBook::GeneratedAttributeMethods
   sig { returns(T::Boolean) }
   def name?; end
 
+  sig { returns(String) }
+  def title; end
+
+  sig { params(value: T.any(String, Symbol)).void }
+  def title=(value); end
+
+  sig { returns(T::Boolean) }
+  def title?; end
+
   sig { returns(Integer) }
   def wizard_id; end
 

diff --git a/spec/test_data/v6.0/expected_spell_book.rbi b/spec/test_data/v6.0/expected_spell_book.rbi
@@ -37,6 +37,15 @@ module SpellBook::GeneratedAttributeMethods
   sig { returns(T::Boolean) }
   def book_type?; end
 
+  sig { returns(String) }
+  def book_category; end
+
+  sig { params(value: T.any(Integer, String, Symbol)).void }
+  def book_category=(value); end
+
+  sig { returns(T::Boolean) }
+  def book_category?; end
+
   sig { returns(Integer) }
   def id; end
 
@@ -55,6 +64,15 @@ module SpellBook::GeneratedAttributeMethods
   sig { returns(T::Boolean) }
   def name?; end
 
+  sig { returns(String) }
+  def title; end
+
+  sig { params(value: T.any(String, Symbol)).void }
+  def title=(value); end
+
+  sig { returns(T::Boolean) }
+  def title?; end
+
   sig { returns(Integer) }
   def wizard_id; end
 

diff --git a/spec/test_data/v6.1/expected_spell_book.rbi b/spec/test_data/v6.1/expected_spell_book.rbi
@@ -37,6 +37,15 @@ module SpellBook::GeneratedAttributeMethods
   sig { returns(T::Boolean) }
   def book_type?; end
 
+  sig { returns(String) }
+  def book_category; end
+
+  sig { params(value: T.any(Integer, String, Symbol)).void }
+  def book_category=(value); end
+
+  sig { returns(T::Boolean) }
+  def book_category?; end
+
   sig { returns(Integer) }
   def id; end
 
@@ -55,6 +64,15 @@ module SpellBook::GeneratedAttributeMethods
   sig { returns(T::Boolean) }
   def name?; end
 
+  sig { returns(String) }
+  def title; end
+
+  sig { params(value: T.any(String, Symbol)).void }
+  def title=(value); end
+
+  sig { returns(T::Boolean) }
+  def title?; end
+
   sig { returns(Integer) }
   def wizard_id; end
 

diff --git a/spec/test_data/v7.0/expected_spell_book.rbi b/spec/test_data/v7.0/expected_spell_book.rbi
@@ -37,6 +37,15 @@ module SpellBook::GeneratedAttributeMethods
   sig { returns(T::Boolean) }
   def book_type?; end
 
+  sig { returns(String) }
+  def book_category; end
+
+  sig { params(value: T.any(Integer, String, Symbol)).void }
+  def book_category=(value); end
+
+  sig { returns(T::Boolean) }
+  def book_category?; end
+
   sig { returns(Integer) }
   def id; end
 
@@ -55,6 +64,15 @@ module SpellBook::GeneratedAttributeMethods
   sig { returns(T::Boolean) }
   def name?; end
 
+  sig { returns(String) }
+  def title; end
+
+  sig { params(value: T.any(String, Symbol)).void }
+  def title=(value); end
+
+  sig { returns(T::Boolean) }
+  def title?; end
+
   sig { returns(Integer) }
   def wizard_id; end