lib/mongomapper/associations.rb in fcoury-mongomapper-0.3.5 vs lib/mongomapper/associations.rb in fcoury-mongomapper-0.4.0

@@ -1,68 +1,63 @@
 module MongoMapper
   module Associations
     module ClassMethods
       def belongs_to(association_id, options = {})
-        association = create_association(:belongs_to, association_id, options)
-
-        ref_id = "#{association_id}_id"
-        key ref_id, String
-
-        define_method("#{ref_id}=") do |value|
-          write_attribute(ref_id, value)
-        end
-
-        if options[:polymorphic]
-          ref_type = "#{association_id}_type"
-          key ref_type, String
-
-          define_method("#{ref_type}=") do |value|
-            write_attribute(ref_type, value)
-          end
-        end
-
-        define_association_methods(association)
-
+        create_association(:belongs_to, association_id, options)        
         self
       end
 
       def many(association_id, options = {})
-        association = create_association(:many, association_id, options)
-        define_association_methods(association)
-
+        create_association(:many, association_id, options)        
         self
       end
 
       def associations
         @associations ||= HashWithIndifferentAccess.new
       end
 
       private
-      def create_association(type, name, options)
-        association = Associations::Base.new(type, name, options)
-        associations[association.name] = association
-        association
-      end
-
-      def define_association_methods(association)
-        define_method(association.name) do
-          get_proxy(association)
+        def create_association(type, name, options)
+          association = Associations::Base.new(type, name, options)
+          associations[association.name] = association
+          define_association_methods(association)
+          define_association_keys(association)
+          association
         end
 
-        define_method("#{association.name}=") do |value|
-          get_proxy(association).replace(value)
-          value
+        def define_association_methods(association)
+          define_method(association.name) do
+            get_proxy(association)
+          end
+          
+          define_method("#{association.name}=") do |value|
+            get_proxy(association).replace(value)
+            value
+          end
         end
-      end
+        
+        def define_association_keys(association)
+          if association.many?
+            if association.polymorphic?
+              association.klass.send :key, association.type_key_name, String
+            end
+          else
+            key "#{association.name}_id", String
+            
+            if association.polymorphic?
+              key association.type_key_name, String
+            end
+          end
+        end
     end
 
     module InstanceMethods
       def get_proxy(association)
-        proxy = self.instance_variable_get(association.ivar)
-        if proxy.nil?
+        unless proxy = self.instance_variable_get(association.ivar)
           proxy = association.proxy_class.new(self, association)
           self.instance_variable_set(association.ivar, proxy)
         end
+        
         proxy
       end
     end
   end
 end