Ticket #17270: manager_queryset.diff

django/db/models/manager.py

@@ -46 +46 @@
     # Tracks each time a Manager instance is created. Used to retain order.
     creation_counter = 0
 
-    def __init__(self):
+    def __init__(self, queryset_cls=None, queryset_method_names=[]):
         super(Manager, self).__init__()
         self._set_creation_counter()
         self.model = None
         self._inherited = False
         self._db = None
+        self._queryset_cls = queryset_cls or QuerySet
+        self.implement_queryset_methods(*queryset_method_names)
 
+    def implement_queryset_cls(self, cls):
+        self._queryset_cls = cls
+
+    def implement_queryset_methods(self, *method_names):
+        for method_name in method_names:
+            self.implement_queryset_method(method_name)
+
+    def implement_queryset_method(self, name):
+        def queryset_method_deco(*args, **kwarg):
+            return getattr(self.get_query_set(), name)(*args, **kwarg)
+        setattr(self, name, queryset_method_deco)
+        
     def contribute_to_class(self, model, name):
         # TODO: Use weakref because of possible memory leak / circular reference.
         self.model = model
@@ -101 +115 @@
     #######################
 
     def get_empty_query_set(self):
-        return EmptyQuerySet(self.model, using=self._db)
+        return self._queryset_cls(self.model, using=self._db).none()
 
     def get_query_set(self):
         """Returns a new QuerySet object.  Subclasses can override this method
         to easily customize the behavior of the Manager.
         """
-        return QuerySet(self.model, using=self._db)
+        return self._queryset_cls(self.model, using=self._db)
 
     def none(self):
         return self.get_empty_query_set()

tests/modeltests/custom_managers/tests.py

@@ -1 +1 @@
 from __future__ import absolute_import
 
+import pickle,tempfile
+
 from django.test import TestCase
 
 from .models import Person, Book, Car, PersonManager, PublishedBookManager
@@ -9 +11 @@
     def test_manager(self):
         p1 = Person.objects.create(first_name="Bugs", last_name="Bunny", fun=True)
         p2 = Person.objects.create(first_name="Droopy", last_name="Dog", fun=False)
+        p3 = Person.objects.create(first_name="Elmer", last_name="Fudd", fun=False)
 
-        self.assertQuerysetEqual(
-            Person.objects.get_fun_people(), [
-                "Bugs Bunny"
-            ],
-            unicode
-        )
+        for objects in [Person.objects, Person.objects_second, Person.objects_third]:
+            self.assertQuerysetEqual(
+                objects.get_fun_people(), [
+                    "Bugs Bunny"
+                ],
+                unicode
+            )
+            
+            self.assertQuerysetEqual(
+                objects.all().get_fun_people(), [
+                    "Bugs Bunny"
+                ],
+                unicode
+            )
+            
+            self.assertQuerysetEqual(
+                objects.get_fun_people().get_not_fun_people(), [],
+                unicode
+            )
+    
+            self.assertQuerysetEqual(
+                objects.filter(id__in=(p1.id,p2.id)).get_fun_people(), [
+                    "Bugs Bunny"
+                ],
+                unicode
+            )
+            
+            self.assertRaises(AttributeError, lambda: objects.all().get_fun_people_manager)
+            
+            # Test dump object
+            f = tempfile.TemporaryFile(mode='wb')
+            try:
+                pickle.dump(objects.all().get_fun_people(), f, pickle.HIGHEST_PROTOCOL)
+            finally:
+                f.close()
+        
         # The RelatedManager used on the 'books' descriptor extends the default
         # manager
         self.assertTrue(isinstance(p2.books, PublishedBookManager))

tests/modeltests/custom_managers/models.py

@@ -13 +13 @@
 
 # An example of a custom manager called "objects".
 
-class PersonManager(models.Manager):
+class ManagerQuerySetFirstPart(models.query.QuerySet):
     def get_fun_people(self):
         return self.filter(fun=True)
+    
+class ManagerQuerySetSecondPart(models.query.QuerySet):
+    def get_not_fun_people(self):
+        return self.filter(fun=False)
 
+class ManagerQuerySetMaster(ManagerQuerySetFirstPart, ManagerQuerySetSecondPart):
+    pass
+
+class ManagerQuerySet(models.query.QuerySet):
+    def get_fun_people(self):
+        return self.filter(fun=True)
+    
+    def get_not_fun_people(self):
+        return self.filter(fun=False)
+    
+class PersonManager(models.Manager):
+    def __init__(self):
+        super(PersonManager, self).__init__()
+        self.implement_queryset_cls(ManagerQuerySet)
+        self.implement_queryset_methods('get_fun_people', 'get_not_fun_people')
+        
+    def get_fun_people_manager(self):
+        return self.filter(fun=True)
+    
+class PersonManagerSecond(models.Manager):
+    def get_fun_people_manager(self):
+        return self.filter(fun=True)
+
 class Person(models.Model):
     first_name = models.CharField(max_length=30)
     last_name = models.CharField(max_length=30)
     fun = models.BooleanField()
     objects = PersonManager()
-
+    objects_second = PersonManagerSecond(ManagerQuerySet, ['get_fun_people', 'get_not_fun_people'])
+    objects_third = PersonManagerSecond(ManagerQuerySetMaster, ['get_fun_people', 'get_not_fun_people'])
+    
     def __unicode__(self):
         return u"%s %s" % (self.first_name, self.last_name)