Ticket #15011: generic.py

from django.contrib.contenttypes import generic
from django.db.models import signals

class ImprovedGenericForeignKey(generic.GenericForeignKey):
        """
        Corrects the behaviour of GenericForeignKey so even if you firstly
        assign an object to this field and then save this object - its PK 
        still gets saved in the fk_field.

        If you assign a not yet saved object to this field an exception is 
        thrown upon saving the model.
        """
        
        class IncompleteData(Exception):
                message = 'Object assigned to field "%s" doesn\'t have a PK (save it first)!'
                
                def __init__(self, field_name):
                        self.field_name = field_name
                        
                def __str__(self):
                        return self.message % self.field_name
                        
        def contribute_to_class(self, cls, name):
                signals.pre_save.connect(self.instance_pre_save, sender=cls, weak=False)
                super(ImprovedGenericForeignKey, self).contribute_to_class(cls, name)
                        
        def instance_pre_save(self, sender, instance, **kwargs):
                """
                Ensures that if GenericForeignKey has an object assigned
                that the fk_field stores the object's PK.
                """
                
                """ If we already have pk set don't do anything... """
                if getattr(instance, self.fk_field) is not None: return
                        
                value = getattr(instance, self.name)
                
                """ 
                If no objects is assigned then we leave it as it is. If null constraints
                are present they should take care of this, if not, well, it's not my fault;)
                """
                if value is not None:
                        fk = value._get_pk_val()
                        
                        if fk is None:
                                raise self.IncompleteData(self.name)
                
                        setattr(instance, self.fk_field, fk)