-
-
Save pzeinlinger/66e0fa833fa30230b1dc0e7b3a281c8e to your computer and use it in GitHub Desktop.
Strict ForeignKeys
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from __future__ import unicode_literals | |
from django.db import models | |
from django.db.models.fields.related_descriptors import ForwardManyToOneDescriptor # noqa | |
class RelationNotLoaded(Exception): | |
pass | |
class StrictForwardManyToOne(ForwardManyToOneDescriptor): | |
def __get__(self, instance, cls=None): | |
try: | |
return getattr(instance, self.cache_name) | |
except AttributeError: | |
raise RelationNotLoaded( | |
'Relation `{rel}` not loaded. Use `select_related` or ' | |
'`fetch_{rel}`'.format(rel=self.field.name) | |
) | |
def explicit_get(self, instance, cls=None): | |
return super(StrictForwardManyToOne, self).__get__(instance, cls) | |
class StrictForeignKey(models.ForeignKey): | |
def contribute_to_class(self, cls, name, **kwargs): | |
super(StrictForeignKey, self).contribute_to_class(cls, name, **kwargs) | |
# Override the descriptor defined by ForeignObject | |
descriptor = StrictForwardManyToOne(self) | |
setattr(cls, self.name, descriptor) | |
# Add a method so you don't always have to use select_related | |
fetch_name = 'fetch_{rel}'.format(rel=self.name) | |
setattr(cls, fetch_name, lambda inst: descriptor.explicit_get(inst)) | |
# Create your models here. | |
class Author(models.Model): | |
name = models.TextField() | |
class Book(models.Model): | |
title = models.TextField() | |
author = StrictForeignKey(Author, on_delete=models.PROTECT, related_name='books') |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment