""" Contact: Filip Sobalski """ 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)