Django soft-deletion module.

softdeletion.py

# -*- coding: utf-8 -*-

"""
Soft-deletion for django resources.

Originally found at: http://datahackermd.com/2013/django-soft-deletion/
"""

import logging

from django.db import models
from django.db.models.query import QuerySet

from django_pg_current_timestamp import CurrentTimestamp


logger = logging.getLogger(__name__)


class LiveField(models.Field):
    """Similar to a BooleanField, but stores False as NULL."""
    description = 'Soft-deletion status'
    __metaclass__ = models.SubfieldBase

    def __init__(self, **kw):
        """NB: keyword args are ignored and only exist here for South migrations compatibility."""
        super(LiveField, self).__init__(default=True, null=True)

    def get_internal_type(self):
        """Create DB column as though for a NullBooleanField."""
        return 'NullBooleanField'

    def get_prep_value(self, value):
        """Convert in-Python value to value we'll store in DB."""
        if value:
            return True
        return None

    def to_python(self, value):
        """NB: Misleading name, since type coercion also occurs when assigning a value to the field in Python."""
        return bool(value)

    def get_prep_lookup(self, lookup_type, value):
        """Filters with .alive=False won't work, so raise a helpful exception instead."""
        if lookup_type == 'exact' and not value:
            msg = "%(model)s doesn't support filters with %(field)s=False. Use a filter with %(field)s=None or an exclude with %(field)s=True instead."
            raise TypeError(msg % {'model': self.model.__name__, 'field': self.name})

        return super(LiveField, self).get_prep_lookup(lookup_type, value)


# South support.
try:
    from south.modelsinspector import add_introspection_rules
    # !!! IMPORTANT !!! Depending on where this module is located the full import path may need to be prepended.
    add_introspection_rules([], [r'^softdeletion\.LiveField'])
except ImportError:
    pass


class SoftDeletionQuerySet(QuerySet):
    def _live_fields(self):
        live_fields = filter(lambda field: isinstance(field, LiveField), self.model._meta.fields)
        return live_fields

    def _live_field_names(self):
        live_fields = self._live_fields()
        live_field_names = [field.name for field in live_fields]
        return live_field_names

    def _auto_now_fields(self):
        auto_now_fields = filter(lambda field: hasattr(field, 'auto_now') and field.auto_now is True, self.model._meta.fields)
        return auto_now_fields

    def delete(self):
        """Bulk delete bypasses individual objects' delete methods."""
        # Update any DateTime fields which specify auto_now=True.
        kw = dict(map(lambda field: (field.name, CurrentTimestamp()), self._auto_now_fields()))
        for live_field_name in self._live_field_names():
            logger.debug('Discovered LiveField with name={}'.format(live_field_name))
            kw[live_field_name] = None
        return super(SoftDeletionQuerySet, self).update(**kw)

    def hard_delete(self):
        return super(SoftDeletionQuerySet, self).delete()

    def alive(self):
        return self.filter(alive=True)

    def dead(self):
        return self.filter(alive=None)


class SoftDeletionManager(models.Manager):
    def __init__(self, *args, **kw):
        self.restrict_to_alive = True
        super(SoftDeletionManager, self).__init__(*args, **kw)

    def get_query_set(self):
        """For legacy compatibility with Django < 1.6."""
        return self.get_queryset()

    def get_queryset(self, restrict_to_alive=''):
        """
        @param `restrict_to_alive` defaults to an empty string, which is used as the default empty/no-arg value which
            signals to use `self.restrict_to_alive`.
        """
        if restrict_to_alive == '':
            restrict_to_alive = self.restrict_to_alive
        qs = SoftDeletionQuerySet(self.model)
        if restrict_to_alive is True:
            qs = qs.alive()
        elif restrict_to_alive is None:
            qs = qs.dead()
        return qs

    def hard_delete(self):
        return self.get_queryset().hard_delete()

    def alive(self):
        qs = self.get_queryset(True)
        return qs

    def dead(self):
        qs = self.get_queryset(None)
        return qs

    def all_with_deleted(self):
        qs = self.get_queryset(False)
        return qs