Customize collector to support deleting elements

Author: yakky

cpkmodel/collector.py

@@ -0,0 +1,75 @@
+from collections import Counter
+from operator import attrgetter
+
+from django.db import transaction
+from django.db.models import signals
+from django.db.models.deletion import Collector
+
+from . import cpkquery as sql
+
+
+class CpkCollector(Collector):
+    def delete(self):
+        # sort instance collections
+        for model, instances in self.data.items():
+            self.data[model] = sorted(instances, key=attrgetter("pk"))
+
+        # if possible, bring the models in an order suitable for databases that
+        # don't support transactions or cannot defer constraint checks until the
+        # end of a transaction.
+        self.sort()
+        # number of objects deleted for each model label
+        deleted_counter = Counter()
+
+        # Optimize for the case with a single obj and no dependencies
+        if len(self.data) == 1 and len(instances) == 1:
+            instance = list(instances)[0]
+            if self.can_fast_delete(instance):
+                with transaction.mark_for_rollback_on_error(self.using):
+                    count = sql.CPkDeleteQuery(model).delete_batch([instance.pk], self.using)
+                setattr(instance, model._meta.pk.attname, None)
+                return count, {model._meta.label: count}
+
+        with transaction.atomic(using=self.using, savepoint=False):
+            # send pre_delete signals
+            for model, obj in self.instances_with_model():
+                if not model._meta.auto_created:
+                    signals.pre_delete.send(sender=model, instance=obj, using=self.using)
+
+            # fast deletes
+            for qs in self.fast_deletes:
+                count = qs._raw_delete(using=self.using)
+                if count:
+                    deleted_counter[qs.model._meta.label] += count
+
+            # update fields
+            for model, instances_for_fieldvalues in self.field_updates.items():
+                for (field, value), instances in instances_for_fieldvalues.items():
+                    query = sql.CPkUpdateQuery(model)
+                    query.update_batch([obj.pk for obj in instances], {field.name: value}, self.using)
+
+            # reverse instance collections
+            for instances in self.data.values():
+                instances.reverse()
+
+            # delete instances
+            for model, instances in self.data.items():
+                query = sql.CPkDeleteQuery(model)
+                pk_list = [obj.pk for obj in instances]
+                count = query.delete_batch(pk_list, self.using)
+                if count:
+                    deleted_counter[model._meta.label] += count
+
+                if not model._meta.auto_created:
+                    for obj in instances:
+                        signals.post_delete.send(sender=model, instance=obj, using=self.using)
+
+        # update collected instances
+        for instances_for_fieldvalues in self.field_updates.values():
+            for (field, value), instances in instances_for_fieldvalues.items():
+                for obj in instances:
+                    setattr(obj, field.attname, value)
+        for model, instances in self.data.items():
+            for instance in instances:
+                setattr(instance, model._meta.pk.attname, None)
+        return sum(deleted_counter.values()), dict(deleted_counter)

cpkmodel/cpkmodel.py

@@ -1,8 +1,8 @@
 from django.db import router
 from django.db.models import Model
 from django.db.models.base import ModelBase
-from django.db.models.deletion import Collector
 
+from .collector import CpkCollector as Collector
 from .compositekey import CompositeKey
 from .constants import CPK_SEP
 from .cpkquery import CPkQuerySet