cascade_improvements.diff

diff --git a/lib/is_paranoid.rb b/lib/is_paranoid.rb
index cad76db..26624e0 100644
--- a/lib/is_paranoid.rb
+++ b/lib/is_paranoid.rb
@@ -38,13 +38,22 @@ module IsParanoid
     # Use update_all with an exclusive scope to restore undo the soft-delete.
     # This bypasses update-related callbacks.
     #
-    # By default, restores cascade through associations that are
+    # By default, restores cascade through associations that are belongs_to
     # :dependent => :destroy and under is_paranoid. You can prevent restoration
     # of associated models by passing :include_destroyed_dependents => false,
     # for example:
     # Android.restore(:include_destroyed_dependents => false)
+    #
+    # Alternatively you can specify which relationships to restore via :include,
+    # for example:
+    # Android.restore(:include => [:parts, memories])
+    # Please note that specifying :include means you're not using
+    # :include_destroyed_dependents by default, though you can explicitly use
+    # both if you want all has_* relationships and specific belongs_to
+    # relationships, for example
+    # Android.restore(:include => [:home, :planet], :include_destroyed_dependents => true)
     def restore(id, options = {})
-      options.reverse_merge!({:include_destroyed_dependents => true})
+      options.reverse_merge!({:include_destroyed_dependents => true}) unless options[:include]
       with_exclusive_scope do
         update_all(
         "#{destroyed_field} = #{connection.quote(field_not_destroyed)}",
@@ -52,13 +61,19 @@ module IsParanoid
         )
       end
 
-      if options[:include_destroyed_dependents]
-        self.reflect_on_all_associations.each do |association|
-          if association.options[:dependent] == :destroy and association.klass.respond_to?(:restore)
-            association.klass.find_destroyed_only(:all,
-            :conditions => ["#{association.primary_key_name} = ?", id]
-            ).each do |model|
-              model.restore
+      self.reflect_on_all_associations.each do |association|
+        if association.options[:dependent] == :destroy and association.klass.respond_to?(:restore)
+          dependent_relationship = association.macro.to_s =~ /^has/
+          if should_restore?(association.name, dependent_relationship, options)
+            if dependent_relationship
+              restore_related(association.klass, association.primary_key_name, id, options)
+            else
+              restore_related(
+                association.klass,
+                association.klass.primary_key,
+                self.first(id).send(association.primary_key_name),
+                options
+              )
             end
           end
         end
@@ -103,6 +118,21 @@ module IsParanoid
         super(name, *args)
       end
     end
+
+    protected
+
+    def should_restore?(association_name, dependent_relationship, options)
+      ([*options[:include]] || []).include?(association_name) or
+        (options[:include_destroyed_dependents] and dependent_relationship)
+    end
+
+    def restore_related klass, key_name, id, options
+      klass.find_destroyed_only(:all,
+        :conditions => ["#{key_name} = ?", id]
+      ).each do |model|
+        model.restore(options)
+      end
+    end
   end
 
   module InstanceMethods
@@ -166,4 +196,4 @@ module IsParanoid
 
 end
 
-ActiveRecord::Base.send(:extend, IsParanoid)
\ No newline at end of file
+ActiveRecord::Base.send(:extend, IsParanoid)
diff --git a/spec/is_paranoid_spec.rb b/spec/is_paranoid_spec.rb
index 21e606e..0c0b6f8 100644
--- a/spec/is_paranoid_spec.rb
+++ b/spec/is_paranoid_spec.rb
@@ -147,6 +147,15 @@ describe IsParanoid do
         @r2d2.restore(:include_destroyed_dependents => false)
       }.should_not change(Component, :count)
     end
+
+    it "should restore parent and child models specified via :include" do
+      sub_component = SubComponent.create(:name => 'part', :component_id => @r2d2.components.first.id)
+      @r2d2.destroy
+      SubComponent.first(:conditions => {:id => sub_component.id}).should be_nil
+      @r2d2.components.first.restore(:include => [:android, :sub_components])
+      SubComponent.first(:conditions => {:id => sub_component.id}).should_not be_nil
+      Android.find(@r2d2.id).should_not be_nil
+    end
   end
 
   describe 'validations' do
@@ -245,4 +254,4 @@ describe IsParanoid do
     end
 
   end
-end
\ No newline at end of file
+end
diff --git a/spec/models.rb b/spec/models.rb
index 832c028..eef35cd 100644
--- a/spec/models.rb
+++ b/spec/models.rb
@@ -21,6 +21,8 @@ end
 
 class Component < ActiveRecord::Base #:nodoc:
   is_paranoid
+  belongs_to :android, :dependent => :destroy
+  has_many :sub_components, :dependent => :destroy
   NEW_NAME = 'Something Else!'
 
   after_destroy :change_name
@@ -29,6 +31,11 @@ class Component < ActiveRecord::Base #:nodoc:
   end
 end
 
+class SubComponent < ActiveRecord::Base #:nodoc:
+  is_paranoid
+  belongs_to :component, :dependent => :destroy
+end
+
 class Memory < ActiveRecord::Base #:nodoc:
   is_paranoid
   belongs_to :android, :class_name => "Android", :foreign_key => "parent_id"
diff --git a/spec/schema.rb b/spec/schema.rb
index 1116b58..5672bd1 100644
--- a/spec/schema.rb
+++ b/spec/schema.rb
@@ -21,6 +21,12 @@ ActiveRecord::Schema.define(:version => 20090317164830) do
     t.datetime "updated_at"
   end
 
+  create_table "sub_components", :force => true do |t|
+    t.string "name"
+    t.integer "component_id"
+    t.datetime "deleted_at"
+  end
+
   create_table "memories", :force => true do |t|
     t.string   "name"
     t.integer  "parent_id"
@@ -42,4 +48,4 @@ ActiveRecord::Schema.define(:version => 20090317164830) do
     t.string   "name"
     t.boolean  "alive", :default => true
   end
-end
\ No newline at end of file
+end