From 486a579ac7dd584214ac7c1763dfda51bb4494b3 Mon Sep 17 00:00:00 2001
From: Abel Daniel <abeld@hal.elte.hu>
Date: Tue, 12 Oct 2010 11:48:02 +0200
Subject: [PATCH] Make .extra() only disable polymorphic query if an arg that
 is currently not handled is actually used. I.e. if all args are supported,
 then do polymorphic query by default.

---
 polymorphic/query.py | 17 ++++++++++++-----
 polymorphic/tests.py |  6 +++++-
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/polymorphic/query.py b/polymorphic/query.py
index e6a0acf..73c0def 100644
--- a/polymorphic/query.py
+++ b/polymorphic/query.py
@@ -85,11 +85,18 @@ class PolymorphicQuerySet(QuerySet):
         return super(PolymorphicQuerySet, self).aggregate(*args, **kwargs)
 
     def extra(self, *args, **kwargs):
-        """only return polymorphic results if we get "polymorphic=True" as a keyword arg."""
-        #for key in kwargs.keys():
-        #    if key not in ['where','order_by', 'params']:
-        #        assert False,"""django_polymorphic: extras() does not support keyword argument %s.
-        self.polymorphic_disabled = not bool(kwargs.pop('polymorphic', False))
+        #
+        # Disable polymorphic, but only if arg that is currently not handled
+        # polymorphically is passed in.
+        #
+        # Allow this to be overridden with 'polymorphic' arg
+        #
+        polymorphic_by_default = True
+        for key in kwargs.keys():
+            if key not in ['where','order_by', 'polymorphic']: # FIXME: also add 'params' here?
+                polymorphic_by_default = False
+        
+        self.polymorphic_disabled = not bool(kwargs.pop('polymorphic', polymorphic_by_default))
         return super(PolymorphicQuerySet, self).extra(*args, **kwargs)
 
     def get_real_instances(self, base_result_objects):
diff --git a/polymorphic/tests.py b/polymorphic/tests.py
index bce835c..f3a69c6 100644
--- a/polymorphic/tests.py
+++ b/polymorphic/tests.py
@@ -357,10 +357,14 @@ __test__ = {"doctest": """
 ### extra() method
 
 >>> Model2A.objects.extra(where=['id IN (2, 3)'])
+[ <Model2B: id 2, field1 (CharField), field2 (CharField)>,
+  <Model2C: id 3, field1 (CharField), field2 (CharField), field3 (CharField)> ]
+
+>>> Model2A.objects.extra(where=['id IN (2, 3)'], select={'dummy_attribute':'1'})
 [ <Model2A: id 2, field1 (CharField)>,
   <Model2A: id 3, field1 (CharField)> ]
 
->>> Model2A.objects.extra(polymorphic=True, where=['id IN (2, 3)'])
+>>> Model2A.objects.extra(polymorphic=True, where=['id IN (2, 3)'], select={'dummy_attribute':'1'})
 [ <Model2B: id 2, field1 (CharField), field2 (CharField)>,
   <Model2C: id 3, field1 (CharField), field2 (CharField), field3 (CharField)> ]