Fix for the bad behaviour of GenericForeignKey field

""" Contact: Filip Sobalski <[email protected]> """

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)