From 07893e6d20ca7b104b5543b8ed5863afff06ff10 Mon Sep 17 00:00:00 2001
From: Simon Charette <simon.charette@zapier.com>
Date: Fri, 4 Oct 2019 23:39:49 +0200
Subject: [PATCH] Added tested support delete(keep_parents).

When a child class is deleted but its ancestor chain is preserved the content
type field should be repointed to the direct parent.
---
 polymodels/models.py | 22 +++++++++++++++++++++-
 tests/test_models.py |  6 ++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/polymodels/models.py b/polymodels/models.py
index 2e26c43..2ada395 100644
--- a/polymodels/models.py
+++ b/polymodels/models.py
@@ -6,7 +6,7 @@
 
 from django.contrib.contenttypes.models import ContentType
 from django.core import checks
-from django.db import models
+from django.db import models, transaction
 from django.db.models.constants import LOOKUP_SEP
 from django.db.models.fields import FieldDoesNotExist
 from django.db.models.signals import class_prepared
@@ -117,6 +117,26 @@ def save(self, *args, **kwargs):
             setattr(self, self.CONTENT_TYPE_FIELD, content_type)
         return super(BasePolymorphicModel, self).save(*args, **kwargs)
 
+    def delete(self, using=None, keep_parents=False):
+        kept_parent = None
+        if keep_parents:
+            parent_ptr = next(iter(self._meta.concrete_model._meta.parents.values()), None)
+            if parent_ptr:
+                kept_parent = getattr(self, parent_ptr.name)
+        if kept_parent:
+            context_manager = transaction.atomic(using=using, savepoint=False)
+        else:
+            context_manager = transaction.mark_for_rollback_on_error(using=using)
+        with context_manager:
+            deletion = super(BasePolymorphicModel, self).delete(
+                using=using, keep_parents=keep_parents
+            )
+            if kept_parent:
+                parent_content_type = get_content_type(kept_parent)
+                setattr(kept_parent, self.CONTENT_TYPE_FIELD, parent_content_type)
+                kept_parent.save(update_fields=[self.CONTENT_TYPE_FIELD])
+        return deletion
+
     @classmethod
     def content_type_lookup(cls, *models, **kwargs):
         query_name = kwargs.pop('query_name', None) or cls.CONTENT_TYPE_FIELD
diff --git a/tests/test_models.py b/tests/test_models.py
index 7e5dfc0..c9ba90e 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -140,6 +140,12 @@ def test_content_type_saving(self):
         self.assertEqual(explicit_mammal.content_type, mammal_content_type)
         self.assertEqual(beaver.content_type, mammal_content_type)
 
+    def test_delete_keep_parents(self):
+        snake = HugeSnake.objects.create(name='snek', length=30)
+        animal_content_type = ContentType.objects.get_for_model(Animal)
+        snake.delete(keep_parents=True)
+        self.assertEqual(snake.animal_ptr.content_type, animal_content_type)
+
 
 class SubclassAccessorsTests(SimpleTestCase):
     def test_dynamic_model_creation_cache_busting(self):