From 9e7a78a8cb871a40064aa6f6ba83860cac75d036 Mon Sep 17 00:00:00 2001 From: Bert Constantin Date: Tue, 2 Feb 2010 07:33:27 +0100 Subject: [PATCH] queryset annotate() method implemented, and testcase --- DOCS.rst | 8 +++++--- polymorphic/polymorphic.py | 31 +++++++++++++++++++++++-------- polymorphic/tests.py | 22 +++++++++++++++++++++- 3 files changed, 49 insertions(+), 12 deletions(-) diff --git a/DOCS.rst b/DOCS.rst index 45f5a12..35cabd9 100644 --- a/DOCS.rst +++ b/DOCS.rst @@ -197,6 +197,11 @@ manage.py dumpdata with Django's seralisation or fixtures (and all polymorphic models use ContentType). This issue seems to be resolved with Django 1.2 (changeset 11863): http://code.djangoproject.com/ticket/7052 + +More Queryset Methods: annotate(), aggregate(), extra() +------------------------------------------------------- + +TODO: add info Custom Managers, Querysets & Inheritance @@ -358,9 +363,6 @@ Currently Unsupported Queryset Methods way. So it seems an implementation would just fall back to the Django vanilla equivalent. -+ ``annotate()``: The current '_get_real_instances' would need minor - enhancement. - + ``defer()`` and ``only()``: Full support, including slight polymorphism enhancements, seems to be straighforward (depends on '_get_real_instances'). diff --git a/polymorphic/polymorphic.py b/polymorphic/polymorphic.py index 770dbd2..0f5e90f 100644 --- a/polymorphic/polymorphic.py +++ b/polymorphic/polymorphic.py @@ -92,11 +92,18 @@ class PolymorphicQuerySet(QuerySet): additional_args = _translate_polymorphic_filter_definitions_in_kwargs(self.model, kwargs) # filter_field='data' return super(PolymorphicQuerySet, self)._filter_or_exclude(negate, *(list(args) + additional_args), **kwargs) + def annotate(self, *args, **kwargs): + for a in args: + assert not '___' in a.lookup, 'PolymorphicModel: annotate(): ___ model lookup supported for keyword arguments only' + for a in kwargs.values(): + a.lookup = _translate_polymorphic_field_path(self.model, a.lookup) + return super(PolymorphicQuerySet, self).annotate(*args, **kwargs) + # these queryset functions are not yet supported def defer(self, *args, **kwargs): raise NotImplementedError def only(self, *args, **kwargs): raise NotImplementedError def aggregate(self, *args, **kwargs): raise NotImplementedError - def annotate(self, *args, **kwargs): raise NotImplementedError + def _get_real_instances(self, base_result_objects): """ @@ -132,9 +139,12 @@ class PolymorphicQuerySet(QuerySet): # - sort base_result_object ids into idlist_per_model lists, depending on their real class; # - also record the correct result order in "ordered_id_list" # - store objects that already have the correct class into "results" + base_result_objects_by_id = {} self_model_content_type_id = ContentType.objects.get_for_model(self.model).pk for base_object in base_result_objects: ordered_id_list.append(base_object.pk) + base_result_objects_by_id[base_object.pk] = base_object + # this object is not a derived object and already the real instance => store it right away if (base_object.polymorphic_ctype_id == self_model_content_type_id): results[base_object.pk] = base_object @@ -144,14 +154,19 @@ class PolymorphicQuerySet(QuerySet): else: idlist_per_model[base_object.get_real_instance_class()].append(base_object.pk) - # for each model in "idlist_per_model" request its objects (the full model) - # from the db and store them in results[] + # For each model in "idlist_per_model" request its objects (the real model) + # from the db and store them in results[]. + # Then we copy the annotate fields from the base objects to the real objects. + # TODO: defer(), only(): support for these would be around here for modelclass, idlist in idlist_per_model.items(): qs = modelclass.base_objects.filter(id__in=idlist) - # copy select related configuration to new qs - qs.dup_select_related(self) - # TODO: defer(), only() and annotate(): support for these would be around here - for o in qs: results[o.pk] = o + qs.dup_select_related(self) # copy select related configuration to new qs + for o in qs: + if self.query.aggregates: + for anno in self.query.aggregates.keys(): + attr = getattr(base_result_objects_by_id[o.pk], anno) + setattr(o, anno, attr) + results[o.pk] = o # re-create correct order and return result list resultlist = [ results[ordered_id] for ordered_id in ordered_id_list if ordered_id in results ] @@ -302,7 +317,7 @@ def _translate_polymorphic_field_path(queryset_model, field_path): Returns: translated path """ classname, sep, pure_field_path = field_path.partition('___') - assert sep == '___' + if not sep: return field_path if '__' in classname: # the user has app label prepended to class name via __ => use Django's get_model function diff --git a/polymorphic/tests.py b/polymorphic/tests.py index cd31973..e07d90f 100644 --- a/polymorphic/tests.py +++ b/polymorphic/tests.py @@ -89,6 +89,16 @@ class MgrInheritB(MgrInheritA): class MgrInheritC(ShowFieldsAndTypes, MgrInheritB): pass +class BlogBase(ShowFieldsAndTypes, PolymorphicModel): + name = models.CharField(max_length=10) +class BlogA(BlogBase): + pass +class BlogB(BlogBase): + pass +class BlogA_Entry(ShowFieldsAndTypes, PolymorphicModel): + blog = models.ForeignKey(BlogA) + text = models.CharField(max_length=10) + # test bad field name #class TestBadFieldModel(PolymorphicModel): # instance_of = models.CharField(max_length=10) @@ -109,6 +119,16 @@ class testclass(TestCase): print 'DiamondXY fields 2: field_b "%s", field_x "%s", field_y "%s"' % (o.field_b, o.field_x, o.field_y) if o.field_b != 'b': print '# Django model inheritance diamond problem detected' + def test_annotate(self): + from django.db.models import Count + + blog = BlogA.objects.create(name='B1') + entry1 = blog.bloga_entry_set.create(text='bla') + entry2 = BlogA_Entry.objects.create(blog=blog, text='bla2') + + qs = BlogBase.objects.annotate(entrycount=Count('BlogA___bloga_entry')) + assert qs[0].entrycount == 2 + __test__ = {"doctest": """ ####################################################### ### Tests @@ -242,4 +262,4 @@ __test__ = {"doctest": """ >>> settings.DEBUG=False """} - \ No newline at end of file +