queryset aggregate() and extra() methods implemented, testcases, docs

.gitignore

@@ -6,14 +6,14 @@
 .pydevproject
 .settings
 
-mypoly.py
-
 tmp
 libraries-local
 
 pushgit
 pushhg
 pushreg
+pbackup
+mcmd.py
 
 ppreadme.py
 ppdocs.py

DOCS.rst

@@ -170,6 +170,19 @@ ManyToManyField, ForeignKey, OneToOneField
           <ModelB: id 2, field1 (CharField), field2 (CharField)>,
           <ModelC: id 3, field1 (CharField), field2 (CharField), field3 (CharField)> ]
 
+annotate(), aggregate() & extra()
+---------------------------------
+
++	``annotate()`` and ``aggregate()`` work just as usual, with the
+	addition that the ``ModelX___field`` syntax can be used for the
+	keyword arguments (but not for the non-keyword arguments).
+	
++	``extra()`` by default works exactly like the vanilla version,
+	with the resulting queryset not being polymorphic. There is
+	experimental support for polymorphic queries with extra() via
+	the keyword argument ``polymorphic=True`` (then only the
+	``where`` and ``order_by`` arguments of extra() should be used).
+
 Non-Polymorphic Queries
 -----------------------
     
@@ -198,11 +211,6 @@ manage.py dumpdata
     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
 ========================================
@@ -359,16 +367,9 @@ Unsupported Methods, Restrictions & Caveats
 Currently Unsupported Queryset Methods
 --------------------------------------
 
-+   ``aggregate()`` probably makes only sense in a purely non-OO/relational
-    way. So it seems an implementation would just fall back to the
-    Django vanilla equivalent.
-    
 +   ``defer()`` and ``only()``: Full support, including slight polymorphism
     enhancements, seems to be straighforward (depends on '_get_real_instances'). 
 
-+   ``extra()``: Does not really work with the current implementation of 
-    '_get_real_instances'. It's unclear if it should be supported.
-
 +	``select_related()`` works just as usual, but it can not (yet) be used
 	to select relations in derived models
 	(like ``ModelA.objects.select_related('ModelC___fieldxy')`` )

polymorphic/polymorphic.py

@@ -17,9 +17,11 @@ Please see LICENSE and AUTHORS for more information.
 from django.db import models
 from django.db.models.base import ModelBase
 from django.db.models.query import QuerySet
+from django.contrib.contenttypes.models import ContentType
+from django import VERSION as django_VERSION
+
 from collections import defaultdict
 from pprint import pprint
-from django.contrib.contenttypes.models import ContentType
 import sys
 
 # chunk-size: maximum number of objects requested per db-request
@@ -76,6 +78,17 @@ class PolymorphicQuerySet(QuerySet):
     is to be used.
     """
 
+    def __init__(self, *args, **kwargs):
+        "init our queryset object member variables"
+        self.polymorphic_disabled = False
+        super(PolymorphicQuerySet, self).__init__(*args, **kwargs)
+        
+    def _clone(self, *args, **kwargs):
+        "Django's _clone only copies its own variables, so we need to copy ours here"
+        new = super(PolymorphicQuerySet, self)._clone(*args, **kwargs)
+        new.polymorphic_disabled = self.polymorphic_disabled
+        return new
+
     def instance_of(self, *args):
         """Filter the queryset to only include the classes in args (and their subclasses).
         Implementation in _translate_polymorphic_filter_defnition."""
@@ -92,18 +105,35 @@ 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):
+    def _process_aggregate_args(self, args, kwargs):
+        """for aggregate and annotate kwargs: allow ModelX___field syntax for kwargs, forbid it for args.
+        Modifies kwargs if needed (these are Aggregate objects, we translate the lookup member variable)"""
         for a in args:
-            assert not '___' in a.lookup, 'PolymorphicModel: annotate(): ___ model lookup supported for keyword arguments only'
+            assert not '___' in a.lookup, 'PolymorphicModel: annotate()/aggregate(): ___ model lookup supported for keyword arguments only'
         for a in kwargs.values():
             a.lookup = _translate_polymorphic_field_path(self.model, a.lookup)
+
+    def annotate(self, *args, **kwargs):
+        """translate the field paths in the kwargs, then call vanilla annotate.
+        _get_real_instances will do the rest of the job after executing the query."""
+        self._process_aggregate_args(args, kwargs)
         return super(PolymorphicQuerySet, self).annotate(*args, **kwargs)
 
+    def aggregate(self, *args, **kwargs):
+        """translate the field paths in the kwargs, then call vanilla aggregate.
+        We need no polymorphic object retrieval for aggregate => switch it off."""
+        self._process_aggregate_args(args, kwargs)
+        self.polymorphic_disabled = True
+        return super(PolymorphicQuerySet, self).aggregate(*args, **kwargs)
+    
+    def extra(self, *args, **kwargs):
+        self.polymorphic_disabled = not kwargs.get('polymorphic',False)
+        if 'polymorphic' in kwargs: kwargs.pop('polymorphic')
+        return super(PolymorphicQuerySet, self).extra(*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 _get_real_instances(self, base_result_objects):
         """
@@ -126,8 +156,10 @@ class PolymorphicQuerySet(QuerySet):
         
         First, we sort the result objects in base_result_objects for their
         subclass (from o.polymorphic_ctype), and then we execute one db query per
-        subclass of objects. Finally we re-sort the resulting objects into the
-        correct order and return them as a list.
+        subclass of objects. Here, we handle any annotations from annotate().
+        
+        Finally we re-sort the resulting objects into the correct order and
+        return them as a list.
         """
         ordered_id_list = []    # list of ids of result-objects in correct order
         results = {}            # polymorphic dict of result-objects, keyed with their id (no order)
@@ -189,6 +221,11 @@ class PolymorphicQuerySet(QuerySet):
         """
         base_iter = super(PolymorphicQuerySet, self).iterator()
 
+        # disabled => work just like a normal queryset
+        if self.polymorphic_disabled:
+            for o in base_iter: yield o
+            raise StopIteration
+
         while True:
             base_result_objects = []
             reached_end = False
@@ -545,8 +582,6 @@ class PolymorphicModelBase(ModelBase):
 ###################################################################################
 ### PolymorphicModel
 
-from django import VERSION as django_VERSION
-
 class PolymorphicModel(models.Model):
     """
     Abstract base class that provides polymorphic behaviour

polymorphic/tests.py

@@ -119,9 +119,10 @@ 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):
+    def test_annotate_aggregate(self):
         from django.db.models import Count
 
+        BlogA.objects.all().delete()
         blog = BlogA.objects.create(name='B1')
         entry1 = blog.bloga_entry_set.create(text='bla')
         entry2 = BlogA_Entry.objects.create(blog=blog, text='bla2')
@@ -129,6 +130,15 @@ class testclass(TestCase):
         qs = BlogBase.objects.annotate(entrycount=Count('BlogA___bloga_entry'))
         assert qs[0].entrycount == 2
 
+        x = BlogBase.objects.aggregate(entrycount=Count('BlogA___bloga_entry'))
+        assert x['entrycount'] == 2
+        
+    def test_extra(self):
+        Model2A.objects.create(field1='A1')
+        Model2B.objects.create(field1='B1', field2='B2')
+        Model2C.objects.create(field1='C1', field2='C2', field3='C3')
+        
+
 __test__ = {"doctest": """
 #######################################################
 ### Tests
@@ -151,6 +161,16 @@ __test__ = {"doctest": """
 >>> o.get_real_instance()
 <Model2C: id 3, field1 (CharField), field2 (CharField), field3 (CharField)>
 
+### extra() method
+
+>>> Model2A.objects.extra(where=['id IN (2, 3)'])
+[ <Model2A: id 2, field1 (CharField)>,
+  <Model2A: id 3, field1 (CharField)> ]
+
+>>> Model2A.objects.extra(polymorphic=True, where=['id IN (2, 3)'])
+[ <Model2B: id 2, field1 (CharField), field2 (CharField)>,
+  <Model2C: id 3, field1 (CharField), field2 (CharField), field3 (CharField)> ]
+
 ### class filtering, instance_of, not_instance_of
 
 >>> Model2A.objects.instance_of(Model2B)