Fixed #20625 -- Chainable Manager/QuerySet methods.

Additionally this patch solves the orthogonal problem that specialized
`QuerySet` like `ValuesQuerySet` didn't inherit from the current `QuerySet`
type. This wasn't an issue until now because we didn't officially support
custom `QuerySet` but it became necessary with the introduction of this new
feature.

Thanks aaugustin, akaariai, carljm, charettes, mjtamlyn, shaib and timgraham
for the reviews.

4 files changed, 147 insertions, 6 deletions

diff --git a/tests/basic/tests.py b/tests/basic/tests.py
index fb21b11279..9d4490afc6 100644
--- a/tests/basic/tests.py
+++ b/tests/basic/tests.py
@@ -6,6 +6,7 @@ import threading
 from django.core.exceptions import ObjectDoesNotExist, MultipleObjectsReturned
 from django.db import connections, DEFAULT_DB_ALIAS
 from django.db.models.fields import Field, FieldDoesNotExist
+from django.db.models.manager import BaseManager
 from django.db.models.query import QuerySet, EmptyQuerySet, ValuesListQuerySet, MAX_GET_RESULTS
 from django.test import TestCase, TransactionTestCase, skipIfDBFeature, skipUnlessDBFeature
 from django.utils import six
@@ -734,3 +735,57 @@ class ConcurrentSaveTests(TransactionTestCase):
         t.join()
         a.save()
         self.assertEqual(Article.objects.get(pk=a.pk).headline, 'foo')
+
+
+class ManagerTest(TestCase):
+    QUERYSET_PROXY_METHODS = [
+        'none',
+        'count',
+        'dates',
+        'datetimes',
+        'distinct',
+        'extra',
+        'get',
+        'get_or_create',
+        'update_or_create',
+        'create',
+        'bulk_create',
+        'filter',
+        'aggregate',
+        'annotate',
+        'complex_filter',
+        'exclude',
+        'in_bulk',
+        'iterator',
+        'earliest',
+        'latest',
+        'first',
+        'last',
+        'order_by',
+        'select_for_update',
+        'select_related',
+        'prefetch_related',
+        'values',
+        'values_list',
+        'update',
+        'reverse',
+        'defer',
+        'only',
+        'using',
+        'exists',
+        '_update',
+    ]
+
+    def test_manager_methods(self):
+        """
+        This test ensures that the correct set of methods from `QuerySet`
+        are copied onto `Manager`.
+
+        It's particularly useful to prevent accidentally leaking new methods
+        into `Manager`. New `QuerySet` methods that should also be copied onto
+        `Manager` will need to be added to `ManagerTest.QUERYSET_PROXY_METHODS`.
+        """
+        self.assertEqual(
+            sorted(BaseManager._get_queryset_methods(QuerySet).keys()),
+            sorted(self.QUERYSET_PROXY_METHODS),
+        )
diff --git a/tests/custom_managers/models.py b/tests/custom_managers/models.py
index 2f5e62fc7a..44d5eb70da 100644
--- a/tests/custom_managers/models.py
+++ b/tests/custom_managers/models.py
@@ -20,6 +20,49 @@ class PersonManager(models.Manager):
     def get_fun_people(self):
         return self.filter(fun=True)
 
+# An example of a custom manager that sets get_queryset().
+
+class PublishedBookManager(models.Manager):
+    def get_queryset(self):
+        return super(PublishedBookManager, self).get_queryset().filter(is_published=True)
+
+# An example of a custom queryset that copies its methods onto the manager.
+
+class CustomQuerySet(models.QuerySet):
+    def filter(self, *args, **kwargs):
+        queryset = super(CustomQuerySet, self).filter(fun=True)
+        queryset._filter_CustomQuerySet = True
+        return queryset
+
+    def public_method(self, *args, **kwargs):
+        return self.all()
+
+    def _private_method(self, *args, **kwargs):
+        return self.all()
+
+    def optout_public_method(self, *args, **kwargs):
+        return self.all()
+    optout_public_method.queryset_only = True
+
+    def _optin_private_method(self, *args, **kwargs):
+        return self.all()
+    _optin_private_method.queryset_only = False
+
+class BaseCustomManager(models.Manager):
+    def __init__(self, arg):
+        super(BaseCustomManager, self).__init__()
+        self.init_arg = arg
+
+    def filter(self, *args, **kwargs):
+        queryset = super(BaseCustomManager, self).filter(fun=True)
+        queryset._filter_CustomManager = True
+        return queryset
+
+    def manager_only(self):
+        return self.all()
+
+CustomManager = BaseCustomManager.from_queryset(CustomQuerySet)
+
 @python_2_unicode_compatible
 class Person(models.Model):
     first_name = models.CharField(max_length=30)
@@ -27,15 +70,12 @@ class Person(models.Model):
     fun = models.BooleanField()
     objects = PersonManager()
 
+    custom_queryset_default_manager = CustomQuerySet.as_manager()
+    custom_queryset_custom_manager = CustomManager('hello')
+
     def __str__(self):
         return "%s %s" % (self.first_name, self.last_name)
 
-# An example of a custom manager that sets get_queryset().
-
-class PublishedBookManager(models.Manager):
-    def get_queryset(self):
-        return super(PublishedBookManager, self).get_queryset().filter(is_published=True)
-
 @python_2_unicode_compatible
 class Book(models.Model):
     title = models.CharField(max_length=50)
diff --git a/tests/custom_managers/tests.py b/tests/custom_managers/tests.py
index 4fe79fe3fb..7fa58b2a88 100644
--- a/tests/custom_managers/tests.py
+++ b/tests/custom_managers/tests.py
@@ -11,12 +11,54 @@ class CustomManagerTests(TestCase):
         p1 = Person.objects.create(first_name="Bugs", last_name="Bunny", fun=True)
         p2 = Person.objects.create(first_name="Droopy", last_name="Dog", fun=False)
 
+        # Test a custom `Manager` method.
         self.assertQuerysetEqual(
             Person.objects.get_fun_people(), [
                 "Bugs Bunny"
             ],
             six.text_type
         )
+
+        # Test that the methods of a custom `QuerySet` are properly
+        # copied onto the default `Manager`.
+        for manager in ['custom_queryset_default_manager',
+                        'custom_queryset_custom_manager']:
+            manager = getattr(Person, manager)
+
+            # Copy public methods.
+            manager.public_method()
+            # Don't copy private methods.
+            with self.assertRaises(AttributeError):
+               manager._private_method()
+            # Copy methods with `manager=True` even if they are private.
+            manager._optin_private_method()
+            # Don't copy methods with `manager=False` even if they are public.
+            with self.assertRaises(AttributeError):
+               manager.optout_public_method()
+
+            # Test that the overriden method is called.
+            queryset = manager.filter()
+            self.assertQuerysetEqual(queryset, ["Bugs Bunny"], six.text_type)
+            self.assertEqual(queryset._filter_CustomQuerySet, True)
+
+            # Test that specialized querysets inherit from our custom queryset.
+            queryset = manager.values_list('first_name', flat=True).filter()
+            self.assertEqual(list(queryset), [six.text_type("Bugs")])
+            self.assertEqual(queryset._filter_CustomQuerySet, True)
+
+        # Test that the custom manager `__init__()` argument has been set.
+        self.assertEqual(Person.custom_queryset_custom_manager.init_arg, 'hello')
+
+        # Test that the custom manager method is only available on the manager.
+        Person.custom_queryset_custom_manager.manager_only()
+        with self.assertRaises(AttributeError):
+            Person.custom_queryset_custom_manager.all().manager_only()
+
+        # Test that the queryset method doesn't override the custom manager method.
+        queryset = Person.custom_queryset_custom_manager.filter()
+        self.assertQuerysetEqual(queryset, ["Bugs Bunny"], six.text_type)
+        self.assertEqual(queryset._filter_CustomManager, True)
+
         # The RelatedManager used on the 'books' descriptor extends the default
         # manager
         self.assertIsInstance(p2.books, PublishedBookManager)
diff --git a/tests/queryset_pickle/tests.py b/tests/queryset_pickle/tests.py
index b4b540c80d..d2f333a9b3 100644
--- a/tests/queryset_pickle/tests.py
+++ b/tests/queryset_pickle/tests.py
@@ -90,3 +90,7 @@ class PickleabilityTestCase(TestCase):
         reloaded = pickle.loads(dumped)
         self.assertEqual(original, reloaded)
         self.assertIs(reloaded.__class__, dynclass)
+
+    def test_specialized_queryset(self):
+        self.assert_pickles(Happening.objects.values('name'))
+        self.assert_pickles(Happening.objects.values('name').dates('when', 'year'))