lib/acts_as.rb in acts_as-0.0.2 vs lib/acts_as.rb in acts_as-0.0.3

@@ -2,65 +2,67 @@
 
 module ActsAs
   class ActsAs::ActiveRecordOnly < StandardError; end
 
   PREFIX = %w(id created_at updated_at)
-  ACTING_FOR = {}
 
   def self.included(base)
+    raise ActiveRecordOnly unless base < ActiveRecord::Base
     base.extend ClassMethods
   end
 
   private
 
   def acts_as_field_match?(method)
-    @association_match = ACTING_FOR.select do |association, fields|
+    @association_match = self.class.acts_as_fields.select do |association, fields|
       fields.select { |f| method.to_s.include?(f) }.any?
     end.keys.first
     @association_match && send(@association_match).respond_to?(method)
   end
 
   module ClassMethods
-    def acts_as(one_association, prefix: [], whitelist: [])
-      define_method(one_association) do |*args|
-        super(*args) || send("build_#{one_association}", *args)
+    def acts_as(association, type = :belongs_to, with: [], prefix: [], **options)
+      type == :belongs_to ? belongs_to(association, **options) : has_one(association, **options)
+      define_method(association) { |*args| super(*args) || send("build_#{association}", *args) }
+
+      if (association_class = new.send(association).class).table_exists?
+        whitelist_and_delegate_fields(association_class, association, prefix, with)
+        override_method_missing
       end
+    end
 
+    def acts_as_fields
+      @acts_as_fields ||= {}
+    end
 
-      association_class = new.send(one_association).class
-      if association_class.table_exists?
+    private
 
-        whitelist_and_delegate_fields(association_class, one_association, prefix, whitelist)
-
-        define_method :method_missing do |method, *args, &block|
-          if acts_as_field_match?(method) then
-            send(@association_match).send(method, *args, &block)
-          else
-            super(method, *args, &block)
-          end
+    def override_method_missing
+      define_method :method_missing do |method, *args, &block|
+        if acts_as_field_match?(method)
+          send(@association_match).send(method, *args, &block)
+        else
+          super(method, *args, &block)
         end
+      end
 
-        define_method :respond_to? do |method, *args, &block|
-          acts_as_field_match?(method) || super(method, *args, &block)
-        end
+      define_method :respond_to? do |method, *args, &block|
+        acts_as_field_match?(method) || super(method, *args, &block)
       end
     end
 
+    def whitelist_and_delegate_fields(association_class, one_association, prefix, with)
+      association_fields = association_class.columns.map(&:name) - PREFIX - prefix + with
 
-    private
-
-    def whitelist_and_delegate_fields(association_class, one_association, prefix, whitelist)
-      association_fields = association_class.columns.map(&:name) - PREFIX - prefix + whitelist
-
       build_prefix_methods(one_association, prefix)
 
       attr_accessible *association_fields
       attr_accessible *prefix.map { |field| "#{one_association}_#{field}" }
 
       delegate(*(association_fields + association_fields.map { |field| "#{field}=" }), to: one_association)
 
       #TODO: This feels like a weird place to remember delegated fields
-      ACTING_FOR[one_association] = association_fields + prefix
+      acts_as_fields[one_association] = association_fields + prefix
     end
 
     def build_prefix_methods(one_association, prefix)
       prefix.each do |field|
         define_method("#{one_association}_#{field}") do |*args|
\ No newline at end of file