Ticket #3871: r17204-custom-reverse-managers.diff

docs/ref/models/relations.txt

@@ -103 +103 @@
 
         Just like ``remove()``, ``clear()`` is only available on
         :class:`~django.db.models.ForeignKey`\s where ``null=True``.
+
+    .. versionadded:: 1.4
+
+    .. method:: manager(manager)
+
+        Returns a new related manager which uses the named ``manager`` instead
+        of the default manager on the related class to look up related objects::
+
+            class Reporter(models.Model):
+                ...
+
+            class Article(models.Model):
+                reporter = models.ForeignKey(Reporter)
+                ...
+                articles = models.Manager()
+                published_articles = PublishedManager()
+
+        In the above example, ``reporter.article_set`` is a manager for all articles
+        written by ``reporter``, whereas ``reporter.article_set.manager("published_articles")``
+        returns a manager with only published articles written by ``reporter``.
+
+        Related managers returned by ``manager()`` do not provide a ``remove()``
+        method. This is enforced to avoid mistakes: ``remove()`` would only check if
+        the given objects are related at all, but not if they are within the domain
+        of the selected manager. For the same reason, managers returned by ``manager()``
+        for either side of a :class:`~django.db.models.ManyToManyField` relation do not
+        have a ``clear()`` method.

django/db/models/fields/related.py

@@ -9 +9 @@
 from django.db.models.query import QuerySet
 from django.db.models.query_utils import QueryWrapper
 from django.db.models.deletion import CASCADE
+from django.db.models.manager import Manager
 from django.utils.encoding import smart_unicode
 from django.utils.translation import ugettext_lazy as _, string_concat
 from django.utils.functional import curry, cached_property
@@ -439 +440 @@
     def related_manager_cls(self):
         # Dynamically create a class that subclasses the related model's default
         # manager.
-        superclass = self.related.model._default_manager.__class__
+        return self._related_manager_cls(self.related.model._default_manager.__class__)
+
+    def _related_manager_cls(self, superclass, uses_default_manager=True):
         rel_field = self.related.field
         rel_model = self.related.model
         attname = rel_field.rel.get_related_field().attname
 
+        def create_related_manager_cls(manager):
+            return self._related_manager_cls(superclass=manager.__class__, uses_default_manager=False)
+
         class RelatedManager(superclass):
             def __init__(self, instance):
                 super(RelatedManager, self).__init__()
@@ -471 +477 @@
                         False,
                         rel_field.related_query_name())
 
+            def manager(self, manager):
+                """
+                Selects a manager from the list of the model's managers, useful
+                in the case of choosing a manager which is not the default, in a reverse relation.
+                """
+                other_manager = getattr(rel_model, manager, None)
+                if isinstance(other_manager, Manager):
+                    rel_manager_cls = create_related_manager_cls(other_manager)
+                    return rel_manager_cls(self.instance)
+                else:
+                    raise AttributeError("Manager '%s' does not exist" % manager)
+
             def add(self, *objs):
                 for obj in objs:
                     if not isinstance(obj, self.model):
@@ -495 +513 @@
 
             # remove() and clear() are only provided if the ForeignKey can have a value of null.
             if rel_field.null:
-                def remove(self, *objs):
-                    val = getattr(self.instance, attname)
-                    for obj in objs:
-                        # Is obj actually part of this descriptor set?
-                        if getattr(obj, rel_field.attname) == val:
-                            setattr(obj, rel_field.name, None)
-                            obj.save()
-                        else:
-                            raise rel_field.rel.to.DoesNotExist("%r is not related to %r." % (obj, self.instance))
-                remove.alters_data = True
+                # If we are dealing with an explicit manager given through the manager method, we
+                # also do not allow the remove() method. We enforce this to avoid surprises: the
+                # current implementation checks only whether the given objects are related, but
+                # not whether they are in the actual domain of the current manager.
+                if uses_default_manager:
+                    def remove(self, *objs):
+                        val = getattr(self.instance, attname)
+                        for obj in objs:
+                            # Is obj actually part of this descriptor set?
+                            if getattr(obj, rel_field.attname) == val:
+                                setattr(obj, rel_field.name, None)
+                                obj.save()
+                            else:
+                                raise rel_field.rel.to.DoesNotExist("%r is not related to %r." % (obj, self.instance))
+                    remove.alters_data = True
 
+                # This works also for explicit managers, in which case it does the expected thing.
                 def clear(self):
                     self.update(**{rel_field.name: None})
                 clear.alters_data = True
@@ -513 +537 @@
         return RelatedManager
 
 
-def create_many_related_manager(superclass, rel):
+def create_many_related_manager(superclass, rel, uses_default_manager=True):
     """Creates a manager that subclasses 'superclass' (which is a Manager)
     and adds behavior for many-to-many related objects."""
     class ManyRelatedManager(superclass):
@@ -570 +594 @@
                     False,
                     self.prefetch_cache_name)
 
+        def manager(self, manager):
+            """
+            Selects a manager from the list of the model's managers, useful
+            in the case of choosing a manager which is not the default, in a reverse relation.
+            """
+            other_manager = getattr(self.model, manager, None)
+            if isinstance(other_manager, Manager):
+                rel_manager_cls = create_many_related_manager(other_manager.__class__, rel, uses_default_manager=False)
+                return rel_manager_cls(
+                    model=self.model, query_field_name=self.query_field_name, instance=self.instance, symmetrical=self.symmetrical,
+                    source_field_name=self.source_field_name, target_field_name=self.target_field_name, reverse=self.reverse,
+                    through=self.through, prefetch_cache_name=self.prefetch_cache_name)
+            else:
+                raise AttributeError("Manager '%s' does not exist" % manager)
+
         # If the ManyToMany relation has an intermediary model,
         # the add and remove methods do not exist.
         if rel.through._meta.auto_created:
@@ -581 +620 @@
                     self._add_items(self.target_field_name, self.source_field_name, *objs)
             add.alters_data = True
 
-            def remove(self, *objs):
-                self._remove_items(self.source_field_name, self.target_field_name, *objs)
+            # The remove method also does not exist when we are
+            # dealing with an explicit manager defined through the
+            # manager method, to avoid mistakes: there is no explicit
+            # check if the removed objects are actually in the related
+            # manager.
+            if uses_default_manager:
+                def remove(self, *objs):
+                    self._remove_items(self.source_field_name, self.target_field_name, *objs)
 
-                # If this is a symmetrical m2m relation to self, remove the mirror entry in the m2m table
-                if self.symmetrical:
-                    self._remove_items(self.target_field_name, self.source_field_name, *objs)
-            remove.alters_data = True
+                    # If this is a symmetrical m2m relation to self, remove the mirror entry in the m2m table
+                    if self.symmetrical:
+                        self._remove_items(self.target_field_name, self.source_field_name, *objs)
+                remove.alters_data = True
 
-        def clear(self):
-            self._clear_items(self.source_field_name)
+        # If an explicit manager has been defined through the manager method,
+        # the clear method is unavailable for the same reason as the remove method.
+        if uses_default_manager:
+            def clear(self):
+                self._clear_items(self.source_field_name)
 
-            # If this is a symmetrical m2m relation to self, clear the mirror entry in the m2m table
-            if self.symmetrical:
-                self._clear_items(self.target_field_name)
-        clear.alters_data = True
+                # If this is a symmetrical m2m relation to self, clear the mirror entry in the m2m table
+                if self.symmetrical:
+                    self._clear_items(self.target_field_name)
+            clear.alters_data = True
 
         def create(self, **kwargs):
             # This check needs to be done here, since we can't later remove this

tests/modeltests/custom_managers/tests.py

@@ -2 +2 @@
 
 from django.test import TestCase
 
-from .models import Person, Book, Car, PersonManager, PublishedBookManager
+from .models import Person, Book, Manufacturer, Color, Car, PersonManager, PublishedBookManager
 
 
 class CustomManagerTests(TestCase):
@@ -42 +42 @@
             lambda b: b.title
         )
 
-        c1 = Car.cars.create(name="Corvette", mileage=21, top_speed=180)
-        c2 = Car.cars.create(name="Neon", mileage=31, top_speed=100)
+    def test_related_manager(self):
+        chevy = Manufacturer.objects.create(name="Chevrolet", country="USA")
+        dodge = Manufacturer.objects.create(name="Dodge", country="USA")
 
+        red = Color.objects.create(name="Red", surcharge=False)
+        white = Color.objects.create(name="White", surcharge=False)
+        silver = Color.objects.create(name="Silver", surcharge=True)
+
+        chevy_corvette = Car.cars.create(name="Corvette", mileage=21, top_speed=180, manufacturer=chevy)
+        chevy_corvette.available_colors.add(red)
+        chevy_corvette.available_colors.add(white)
+
+        dodge_neon = Car.cars.create(name="Neon", mileage=31, top_speed=100, manufacturer=dodge)
+        dodge_neon.available_colors.add(white)
+        dodge_neon.available_colors.add(silver)
+
+        dodge_viper = Car.cars.create(name="Viper", mileage=14, top_speed=200, manufacturer=dodge)
+        dodge_viper.available_colors.add(red)
+
         self.assertQuerysetEqual(
             Car.cars.order_by("name"), [
                 "Corvette",
                 "Neon",
+                "Viper",
             ],
             lambda c: c.name
         )
 
         self.assertQuerysetEqual(
-            Car.fast_cars.all(), [
+            Car.fast_cars.order_by("name"), [
                 "Corvette",
+                "Viper",
             ],
             lambda c: c.name
         )
 
+        self.assertQuerysetEqual(
+            dodge.car_set.order_by("name"), [
+                "Neon",
+                "Viper",
+            ],
+            lambda c: c.name
+        )
+
+        self.assertQuerysetEqual(
+            white.car_set.order_by("name"), [
+                "Corvette",
+                "Neon",
+            ],
+            lambda c: c.name
+        )
+
+        self.assertQuerysetEqual(
+            dodge_neon.available_colors.order_by("name"), [
+                "Silver",
+                "White",
+            ],
+            lambda c: c.name
+        )
+
+        # The remove and clear methods must be available. Check only
+        # for availability here but don't actually run those methods
+        # -> no "()".
+        dodge.car_set.remove
+        dodge.car_set.clear
+        white.car_set.remove
+        white.car_set.clear
+        dodge_neon.available_colors.remove
+        dodge_neon.available_colors.clear
+
         # Each model class gets a "_default_manager" attribute, which is a
         # reference to the first manager defined in the class. In this case,
         # it's "cars".
-
         self.assertQuerysetEqual(
             Car._default_manager.order_by("name"), [
                 "Corvette",
                 "Neon",
+                "Viper",
             ],
             lambda c: c.name
         )
+
+        # Choosing a custom manager in a reverse relation.
+        self.assertQuerysetEqual(
+            dodge.car_set.manager("fast_cars").all(), [
+                "Viper",
+            ],
+            lambda c: c.name
+        )
+
+        # Choosing a custom manager in a many-to-many reverse relation.
+        self.assertQuerysetEqual(
+            white.car_set.manager("fast_cars").all(), [
+                "Corvette",
+            ],
+            lambda c: c.name
+        )
+
+        # Choosing a custom manager in many-to-many forward relation.
+        self.assertQuerysetEqual(
+            dodge_neon.available_colors.manager("free_of_charge_colors").all(), [
+                "White",
+            ],
+            lambda c: c.name
+        )
+
+        # Unknown managers must raise exception.
+        self.assertRaises(
+            AttributeError,
+            dodge.car_set.manager, "red_cars"
+        )
+        self.assertRaises(
+            AttributeError,
+            white.car_set.manager, "red_cars"
+        )
+        self.assertRaises(
+            AttributeError,
+            dodge_neon.available_colors.manager, "beautiful_colors"
+        )
+
+        # Removing from managers with explicitly selected related manager
+        # must not be possible to avoid mistakes (as the related manager
+        # does not check if the given objects are actually part of the
+        # domain selected by that manager).
+        # clear on (regular) reverse relations is the only exception.
+        self.assertRaises(AttributeError, lambda: dodge.car_set.manager("fast_cars").remove)
+        dodge.car_set.manager("fast_cars").clear  # check availability, but don't run method
+        self.assertRaises(AttributeError, lambda: white.car_set.manager("fast_cars").remove)
+        self.assertRaises(AttributeError, lambda: white.car_set.manager("fast_cars").clear)
+        self.assertRaises(AttributeError, lambda: dodge_neon.available_colors.manager("free_of_charge_colors").remove)
+        self.assertRaises(AttributeError, lambda: dodge_neon.available_colors.manager("free_of_charge_colors").clear)
+
+        # Clearing reverse relation with custom manager must only
+        # remove objects selected by that manager.
+        self.assertEqual(dodge.car_set.count(), 2)
+        dodge.car_set.manager("fast_cars").clear()
+        self.assertEqual(dodge.car_set.count(), 1)
+        dodge.car_set.clear()
+        self.assertEqual(dodge.car_set.count(), 0)

tests/modeltests/custom_managers/models.py

@@ -44 +44 @@
 
 # An example of providing multiple custom managers.
 
+class Manufacturer(models.Model):
+    name = models.CharField(max_length=10)
+    country = models.CharField(max_length=20)
+
+class FreeOfChargeColorManager(models.Manager):
+    def get_query_set(self):
+        return super(FreeOfChargeColorManager, self).get_query_set().filter(surcharge=False)
+
+class Color(models.Model):
+    name = models.CharField(max_length=10)
+    surcharge = models.BooleanField()
+    objects = models.Manager()
+    free_of_charge_colors = FreeOfChargeColorManager()
+
 class FastCarManager(models.Manager):
     def get_query_set(self):
         return super(FastCarManager, self).get_query_set().filter(top_speed__gt=150)
 
 class Car(models.Model):
     name = models.CharField(max_length=10)
+    manufacturer = models.ForeignKey(Manufacturer, null=True)
+    available_colors = models.ManyToManyField(Color)
     mileage = models.IntegerField()
     top_speed = models.IntegerField(help_text="In miles per hour.")
     cars = models.Manager()