from django.db import models
from django.core.exceptions import ObjectDoesNotExist


class OneToOneField(models.OneToOneField):
    def __init__(self, to, **kwargs):
        self.related_default = kwargs.pop('related_default', None)
        super(OneToOneField, self).__init__(to, **kwargs)

    def contribute_to_related_class(self, cls, related):
        setattr(cls, related.get_accessor_name(),
                SingleRelatedObjectDescriptor(related, self.related_default))


class SingleRelatedObjectDescriptor(models.related.SingleRelatedObjectDescriptor):
    def __init__(self, related, default):
        super(SingleRelatedObjectDescriptor, self).__init__(related)
        self.default = default

    def __get__(self, instance, instance_type=None):
        try:
            return super(SingleRelatedObjectDescriptor, self).__get__(instance,
                                                                      instance_type)
        except ObjectDoesNotExist:
            if self.default is None:
                raise
            value = self.default(instance)
            self.__set__(instance, value)
            return value