Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 21 additions & 26 deletions polymodels/fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from django.contrib.contenttypes.models import ContentType
from django.db.models import ForeignKey, Q
from django.db.models.fields import NOT_PROVIDED
from django.db.models.fields.related import ManyToOneRel, RelatedField
from django.db.models.fields.related import RelatedField
from django.utils.deconstruct import deconstructible
from django.utils.functional import LazyObject
from django.utils.six import string_types
Expand All @@ -18,28 +18,26 @@
from .utils import get_content_type


class PolymorphicManyToOneRel(ManyToOneRel):
"""
Relationship that generates a `limit_choices_to` based on it's polymorphic type subclasses.
"""
class LimitChoicesToSubclasses(object):
def __init__(self, field, limit_choices_to):
self.field = field
self.limit_choices_to = limit_choices_to

@property
def limit_choices_to(self):
subclasses_lookup = self.polymorphic_type.subclasses_lookup('pk')
limit_choices_to = self._limit_choices_to
def value(self):
subclasses_lookup = self.field.polymorphic_type.subclasses_lookup('pk')
limit_choices_to = self.limit_choices_to
if limit_choices_to is None:
limit_choices_to = subclasses_lookup
limit_choices_to = subclasses_lookup.copy()
elif isinstance(limit_choices_to, dict):
limit_choices_to = dict(limit_choices_to, **subclasses_lookup)
elif isinstance(limit_choices_to, Q):
limit_choices_to = limit_choices_to & Q(**subclasses_lookup)
self.__dict__['limit_choices_to'] = limit_choices_to
self.__dict__['value'] = limit_choices_to
return limit_choices_to

@limit_choices_to.setter
def limit_choices_to(self, value):
self._limit_choices_to = value
self.__dict__.pop('limit_choices_to', None)
def __call__(self):
return self.value


class LazyPolymorphicTypeQueryset(LazyObject):
Expand All @@ -51,7 +49,7 @@ def _setup(self):
remote_field = self.__dict__.get('remote_field')
db = self.__dict__.get('db')
self._wrapped = get_remote_model(remote_field)._default_manager.using(db).complex_filter(
remote_field.limit_choices_to
remote_field.limit_choices_to()
)


Expand Down Expand Up @@ -81,21 +79,19 @@ class PolymorphicTypeField(ForeignKey):
description = _(
'Content type of a subclass of %(type)s'
)
rel_class = PolymorphicManyToOneRel

def __init__(self, polymorphic_type, *args, **kwargs):
if not isinstance(polymorphic_type, string_types):
self.validate_polymorphic_type(polymorphic_type)
self.polymorphic_type = polymorphic_type
limit_choices_to = LimitChoicesToSubclasses(self, kwargs.pop('limit_choices_to', None))
defaults = {
'to': ContentType,
'related_name': '+',
'limit_choices_to': limit_choices_to,
}
# TODO: Remove when support for Django 1.8 is dropped.
if not hasattr(ForeignKey, 'rel_class'):
defaults['rel_class'] = PolymorphicManyToOneRel
defaults.update(kwargs)
defaults.update(**kwargs)
super(PolymorphicTypeField, self).__init__(*args, **defaults)
get_remote_field(self).polymorphic_type = polymorphic_type

def validate_polymorphic_type(self, model):
if not isclass(model) or not issubclass(model, BasePolymorphicModel):
Expand All @@ -106,12 +102,11 @@ def validate_polymorphic_type(self, model):

def contribute_to_class(self, cls, name):
super(PolymorphicTypeField, self).contribute_to_class(cls, name)
polymorphic_type = get_remote_field(self).polymorphic_type
polymorphic_type = self.polymorphic_type
if (isinstance(polymorphic_type, string_types) or
polymorphic_type._meta.pk is None):
def resolve_polymorphic_type(model, related_model, field):
field.validate_polymorphic_type(related_model)
get_remote_field(field).polymorphic_type = related_model
field.do_polymorphic_type(related_model)
lazy_related_operation(resolve_polymorphic_type, cls, polymorphic_type, field=self)
else:
Expand All @@ -121,18 +116,18 @@ def do_polymorphic_type(self, polymorphic_type):
if self.default is NOT_PROVIDED and not self.null:
opts = polymorphic_type._meta
self.default = ContentTypeReference(opts.app_label, opts.model_name)
self.polymorphic_type = polymorphic_type
self.type = polymorphic_type.__name__
self.error_messages['invalid'] = (
'Specified content type is not of a subclass of %s.' % polymorphic_type._meta.object_name
)

def formfield(self, **kwargs):
db = kwargs.pop('using', None)
remote_model = get_remote_model(get_remote_field(self))
if isinstance(remote_model, string_types):
if isinstance(self.polymorphic_type, string_types):
raise ValueError(
"Cannot create form field for %r yet, because its related model %r has not been loaded yet" % (
self.name, remote_model
self.name, self.polymorphic_type
)
)
defaults = {
Expand Down
43 changes: 13 additions & 30 deletions tests/test_fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,41 +43,24 @@ def test_limit_choices_to(self):
"""
Make sure existing `limit_choices_to` are taken into consideration
"""
type_field = Trait._meta.get_field('trait_type')
remote_field = get_remote_field(type_field)

# Make sure it's cached
limit_choices_to = remote_field.limit_choices_to
self.assertIn('limit_choices_to', remote_field.__dict__)

extra_limit_choices_to = {'app_label': 'polymodels'}

# Make sure it works with existing dict `limit_choices_to`
remote_field.limit_choices_to = extra_limit_choices_to
# Cache should be cleared
self.assertNotIn('limit_choices_to', remote_field.__dict__)
field = PolymorphicTypeField(Trait, on_delete=models.CASCADE)
remote_field = get_remote_field(field)
subclasses_lookup = Trait.subclasses_lookup('pk')
self.assertEqual(remote_field.limit_choices_to(), subclasses_lookup)
# Test dict() limit_choices_to.
limit_choices_to = {'app_label': 'polymodels'}
field = PolymorphicTypeField(Trait, on_delete=models.CASCADE, limit_choices_to=limit_choices_to)
remote_field = get_remote_field(field)
self.assertEqual(
remote_field.limit_choices_to,
dict(extra_limit_choices_to, **limit_choices_to)
remote_field.limit_choices_to(), dict(subclasses_lookup, **limit_choices_to)
)

# Make sure it works with existing Q `limit_choices_to`
remote_field.limit_choices_to = Q(**extra_limit_choices_to)
# Cache should be cleared
self.assertNotIn('limit_choices_to', remote_field.__dict__)
remote_field_limit_choices_to = remote_field.limit_choices_to
self.assertEqual(remote_field_limit_choices_to.connector, Q.AND)
self.assertFalse(remote_field_limit_choices_to.negated)
# Test Q() limit_choices_to.
field = PolymorphicTypeField(Trait, on_delete=models.CASCADE, limit_choices_to=Q(**limit_choices_to))
remote_field = get_remote_field(field)
self.assertEqual(
remote_field_limit_choices_to.children,
list(extra_limit_choices_to.items()) + list(limit_choices_to.items())
str(remote_field.limit_choices_to()), str(Q(**limit_choices_to) & Q(**subclasses_lookup))
)

# Re-assign the original value
remote_field.limit_choices_to = None
# Cache should be cleared
self.assertNotIn('limit_choices_to', remote_field.__dict__)

def test_invalid_type(self):
trait = Trait.objects.create()
snake_type = get_content_type(Snake)
Expand Down