Customizable STI rule for ActiveRecord

active_sti.rb

# usage:
#  class User < AR::Base
#    extend ActiveSTI
#    define_sti_rule do|record|
#      if record['parent']
#        ChildUser
#      else
#        RootUser
#      end
#    end
#  end
#
# tested with activerecord 3.2.11

module ActiveSTI
  def instantiate(record)
    sti_class = find_active_sti_class(record)
    record_id = sti_class.primary_key && record[sti_class.primary_key]

    if ActiveRecord::IdentityMap.enabled? && record_id
      instance = use_identity_map(sti_class, record_id, record)
    else
      instance = sti_class.allocate.init_with('attributes' => record)
    end

    instance
  end

  def find_active_sti_class(record)
    if @@sti_rule
      @@sti_rule[record]
    else
      find_sti_class(record[inheritance_column])
    end
  end

  # (record -> sti_class) ->
  def define_sti_rule(&block)
    @@sti_rule = block
  end
end