module Enumerize::ActiveRecordSupport::InstanceMethods

def becomes(klass)
  became = super
  klass.enumerized_attributes.each do |attr|
    # Rescue when column associated to the enum does not exist.
    begin
      became.send("#{attr.name}=", send(attr.name))
    rescue ActiveModel::MissingAttributeError
    end
  end
  became
end

def reload(options = nil)
  reloaded = super
  reloaded.class.enumerized_attributes.each do |attr|
    begin
      # Checks first if the enumerized attribute is in ActiveRecord::Store
      store_attr, _ = reloaded.class.stored_attributes.detect do |_store_attr, keys|
        keys.include?(attr.name)
      end
      if store_attr.present?
        unless reloaded.send(store_attr).nil?
          reloaded.send("#{attr.name}=", reloaded.send(store_attr).with_indifferent_access[attr.name])
        end
      else
        reloaded.send("#{attr.name}=", reloaded[attr.name])
      end
    rescue ActiveModel::MissingAttributeError
    end
  end
  reloaded
end

def write_attribute(attr_name, value, *options)
  if self.class.enumerized_attributes[attr_name]
    _enumerized_values_for_validation[attr_name.to_s] = value
  end
  super
end