InheritanceQuerySet annotate patch

# HG changeset patch
# User Facundo Gaich <facugaich@gmail.com>
# Date 1325011816 10800
# Node ID d13f6f7fe1bc3ac4cc108f758dbc8d43b50a13a6
# Parent  8113fd6b7bf64ce35fe254bc5a816d792ae9c654
Change _clone logic to keep annotations done before select_subclasses

diff -r 8113fd6b7bf6 -r d13f6f7fe1bc model_utils/managers.py
--- a/model_utils/managers.py   Fri Dec 23 21:03:04 2011 -0500
+++ b/model_utils/managers.py   Tue Dec 27 15:50:16 2011 -0300
@@ -8,10 +8,6 @@
 from django.db.models.query import QuerySet

 class InheritanceQuerySet(QuerySet):
-    def __init__(self, *args, **kwargs):
-        self._annotated = None
-        super(InheritanceQuerySet, self).__init__(*args, **kwargs)
-
     def select_subclasses(self, *subclasses):
         if not subclasses:
             subclasses = [o for o in dir(self.model)
@@ -22,11 +18,9 @@
         return new_qs

     def _clone(self, klass=None, setup=False, **kwargs):
-        try:
-            kwargs.update({'subclasses': self.subclasses,
-                           '_annotated': self._annotated})
-        except AttributeError:
-            pass
+        for name in ['subclasses', '_annotated']:
+            if hasattr(self, name):
+                kwargs[name] = getattr(self, name)
         return super(InheritanceQuerySet, self)._clone(klass, setup, **kwargs)

     def annotate(self, *args, **kwargs):
@@ -40,7 +34,7 @@
             for obj in iter:
                 sub_obj = [getattr(obj, s) for s in self.subclasses if getattr(obj, s)] or [obj]
                 sub_obj = sub_obj[0]
-                if self._annotated:
+                if getattr(self, '_annotated', False):
                     for k in self._annotated:
                         setattr(sub_obj, k, getattr(obj, k))

diff -r 8113fd6b7bf6 -r d13f6f7fe1bc model_utils/tests/tests.py
--- a/model_utils/tests/tests.py    Fri Dec 23 21:03:04 2011 -0500
+++ b/model_utils/tests/tests.py    Tue Dec 27 15:50:16 2011 -0300
@@ -281,6 +281,14 @@
         qs = InheritanceManagerTestParent.objects.select_subclasses().annotate(test_count=models.Count('id'))
         self.assertEqual(qs.get(id=self.child1.id).test_count, 1)

+    def test_annotate_before_select_subclasses(self):
+        qs = InheritanceManagerTestParent.objects.annotate(models.Count('id')).select_subclasses()
+        self.assertEqual(qs.get(id=self.child1.id).id__count, 1)
+
+    def test_annotate_with_named_arguments_before_select_subclasses(self):
+        qs = InheritanceManagerTestParent.objects.annotate(test_count=models.Count('id')).select_subclasses()
+        self.assertEqual(qs.get(id=self.child1.id).test_count, 1)
+

 class TimeStampedModelTests(TestCase):
     def test_created(self):