lib/polyamorous/join_association.rb in polyamorous-0.6.3 vs lib/polyamorous/join_association.rb in polyamorous-0.6.4

@@ -20,22 +20,22 @@
 
     def initialize_with_polymorphism(reflection, join_dependency, parent = nil, polymorphic_class = nil)
       if polymorphic_class && ::ActiveRecord::Base > polymorphic_class
         swapping_reflection_klass(reflection, polymorphic_class) do |reflection|
           initialize_without_polymorphism(reflection, join_dependency, parent)
+          self.reflection.options[:polymorphic] = true
         end
       else
         initialize_without_polymorphism(reflection, join_dependency, parent)
       end
     end
 
     def swapping_reflection_klass(reflection, klass)
-      reflection = reflection.clone
-      original_polymorphic = reflection.options.delete(:polymorphic)
-      reflection.instance_variable_set(:@klass, klass)
-      yield reflection
-    ensure
-      reflection.options[:polymorphic] = original_polymorphic
+      new_reflection = reflection.clone
+      new_reflection.instance_variable_set(:@options, reflection.options.clone)
+      new_reflection.options.delete(:polymorphic)
+      new_reflection.instance_variable_set(:@klass, klass)
+      yield new_reflection
     end
 
     def equality_with_polymorphism(other)
       equality_without_polymorphism(other) && base_klass == other.base_klass
     end