Code

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

File r17204-custom-reverse-managers.diff, 16.7 KB (added by sebastian, 3 years ago)

Revised patch.

Line 
1Index: docs/ref/models/relations.txt
2===================================================================
3--- docs/ref/models/relations.txt       (revision 17204)
4+++ docs/ref/models/relations.txt       (working copy)
5@@ -103,3 +103,30 @@
6 
7         Just like ``remove()``, ``clear()`` is only available on
8         :class:`~django.db.models.ForeignKey`\s where ``null=True``.
9+
10+    .. versionadded:: 1.4
11+
12+    .. method:: manager(manager)
13+
14+        Returns a new related manager which uses the named ``manager`` instead
15+        of the default manager on the related class to look up related objects::
16+
17+            class Reporter(models.Model):
18+                ...
19+
20+            class Article(models.Model):
21+                reporter = models.ForeignKey(Reporter)
22+                ...
23+                articles = models.Manager()
24+                published_articles = PublishedManager()
25+
26+        In the above example, ``reporter.article_set`` is a manager for all articles
27+        written by ``reporter``, whereas ``reporter.article_set.manager("published_articles")``
28+        returns a manager with only published articles written by ``reporter``.
29+
30+        Related managers returned by ``manager()`` do not provide a ``remove()``
31+        method. This is enforced to avoid mistakes: ``remove()`` would only check if
32+        the given objects are related at all, but not if they are within the domain
33+        of the selected manager. For the same reason, managers returned by ``manager()``
34+        for either side of a :class:`~django.db.models.ManyToManyField` relation do not
35+        have a ``clear()`` method.
36Index: django/db/models/fields/related.py
37===================================================================
38--- django/db/models/fields/related.py  (revision 17204)
39+++ django/db/models/fields/related.py  (working copy)
40@@ -9,6 +9,7 @@
41 from django.db.models.query import QuerySet
42 from django.db.models.query_utils import QueryWrapper
43 from django.db.models.deletion import CASCADE
44+from django.db.models.manager import Manager
45 from django.utils.encoding import smart_unicode
46 from django.utils.translation import ugettext_lazy as _, string_concat
47 from django.utils.functional import curry, cached_property
48@@ -439,11 +440,16 @@
49     def related_manager_cls(self):
50         # Dynamically create a class that subclasses the related model's default
51         # manager.
52-        superclass = self.related.model._default_manager.__class__
53+        return self._related_manager_cls(self.related.model._default_manager.__class__)
54+
55+    def _related_manager_cls(self, superclass, uses_default_manager=True):
56         rel_field = self.related.field
57         rel_model = self.related.model
58         attname = rel_field.rel.get_related_field().attname
59 
60+        def create_related_manager_cls(manager):
61+            return self._related_manager_cls(superclass=manager.__class__, uses_default_manager=False)
62+
63         class RelatedManager(superclass):
64             def __init__(self, instance):
65                 super(RelatedManager, self).__init__()
66@@ -471,6 +477,18 @@
67                         False,
68                         rel_field.related_query_name())
69 
70+            def manager(self, manager):
71+                """
72+                Selects a manager from the list of the model's managers, useful
73+                in the case of choosing a manager which is not the default, in a reverse relation.
74+                """
75+                other_manager = getattr(rel_model, manager, None)
76+                if isinstance(other_manager, Manager):
77+                    rel_manager_cls = create_related_manager_cls(other_manager)
78+                    return rel_manager_cls(self.instance)
79+                else:
80+                    raise AttributeError("Manager '%s' does not exist" % manager)
81+
82             def add(self, *objs):
83                 for obj in objs:
84                     if not isinstance(obj, self.model):
85@@ -495,17 +513,23 @@
86 
87             # remove() and clear() are only provided if the ForeignKey can have a value of null.
88             if rel_field.null:
89-                def remove(self, *objs):
90-                    val = getattr(self.instance, attname)
91-                    for obj in objs:
92-                        # Is obj actually part of this descriptor set?
93-                        if getattr(obj, rel_field.attname) == val:
94-                            setattr(obj, rel_field.name, None)
95-                            obj.save()
96-                        else:
97-                            raise rel_field.rel.to.DoesNotExist("%r is not related to %r." % (obj, self.instance))
98-                remove.alters_data = True
99+                # If we are dealing with an explicit manager given through the manager method, we
100+                # also do not allow the remove() method. We enforce this to avoid surprises: the
101+                # current implementation checks only whether the given objects are related, but
102+                # not whether they are in the actual domain of the current manager.
103+                if uses_default_manager:
104+                    def remove(self, *objs):
105+                        val = getattr(self.instance, attname)
106+                        for obj in objs:
107+                            # Is obj actually part of this descriptor set?
108+                            if getattr(obj, rel_field.attname) == val:
109+                                setattr(obj, rel_field.name, None)
110+                                obj.save()
111+                            else:
112+                                raise rel_field.rel.to.DoesNotExist("%r is not related to %r." % (obj, self.instance))
113+                    remove.alters_data = True
114 
115+                # This works also for explicit managers, in which case it does the expected thing.
116                 def clear(self):
117                     self.update(**{rel_field.name: None})
118                 clear.alters_data = True
119@@ -513,7 +537,7 @@
120         return RelatedManager
121 
122 
123-def create_many_related_manager(superclass, rel):
124+def create_many_related_manager(superclass, rel, uses_default_manager=True):
125     """Creates a manager that subclasses 'superclass' (which is a Manager)
126     and adds behavior for many-to-many related objects."""
127     class ManyRelatedManager(superclass):
128@@ -570,6 +594,21 @@
129                     False,
130                     self.prefetch_cache_name)
131 
132+        def manager(self, manager):
133+            """
134+            Selects a manager from the list of the model's managers, useful
135+            in the case of choosing a manager which is not the default, in a reverse relation.
136+            """
137+            other_manager = getattr(self.model, manager, None)
138+            if isinstance(other_manager, Manager):
139+                rel_manager_cls = create_many_related_manager(other_manager.__class__, rel, uses_default_manager=False)
140+                return rel_manager_cls(
141+                    model=self.model, query_field_name=self.query_field_name, instance=self.instance, symmetrical=self.symmetrical,
142+                    source_field_name=self.source_field_name, target_field_name=self.target_field_name, reverse=self.reverse,
143+                    through=self.through, prefetch_cache_name=self.prefetch_cache_name)
144+            else:
145+                raise AttributeError("Manager '%s' does not exist" % manager)
146+
147         # If the ManyToMany relation has an intermediary model,
148         # the add and remove methods do not exist.
149         if rel.through._meta.auto_created:
150@@ -581,21 +620,30 @@
151                     self._add_items(self.target_field_name, self.source_field_name, *objs)
152             add.alters_data = True
153 
154-            def remove(self, *objs):
155-                self._remove_items(self.source_field_name, self.target_field_name, *objs)
156+            # The remove method also does not exist when we are
157+            # dealing with an explicit manager defined through the
158+            # manager method, to avoid mistakes: there is no explicit
159+            # check if the removed objects are actually in the related
160+            # manager.
161+            if uses_default_manager:
162+                def remove(self, *objs):
163+                    self._remove_items(self.source_field_name, self.target_field_name, *objs)
164 
165-                # If this is a symmetrical m2m relation to self, remove the mirror entry in the m2m table
166-                if self.symmetrical:
167-                    self._remove_items(self.target_field_name, self.source_field_name, *objs)
168-            remove.alters_data = True
169+                    # If this is a symmetrical m2m relation to self, remove the mirror entry in the m2m table
170+                    if self.symmetrical:
171+                        self._remove_items(self.target_field_name, self.source_field_name, *objs)
172+                remove.alters_data = True
173 
174-        def clear(self):
175-            self._clear_items(self.source_field_name)
176+        # If an explicit manager has been defined through the manager method,
177+        # the clear method is unavailable for the same reason as the remove method.
178+        if uses_default_manager:
179+            def clear(self):
180+                self._clear_items(self.source_field_name)
181 
182-            # If this is a symmetrical m2m relation to self, clear the mirror entry in the m2m table
183-            if self.symmetrical:
184-                self._clear_items(self.target_field_name)
185-        clear.alters_data = True
186+                # If this is a symmetrical m2m relation to self, clear the mirror entry in the m2m table
187+                if self.symmetrical:
188+                    self._clear_items(self.target_field_name)
189+            clear.alters_data = True
190 
191         def create(self, **kwargs):
192             # This check needs to be done here, since we can't later remove this
193Index: tests/modeltests/custom_managers/tests.py
194===================================================================
195--- tests/modeltests/custom_managers/tests.py   (revision 17204)
196+++ tests/modeltests/custom_managers/tests.py   (working copy)
197@@ -2,7 +2,7 @@
198 
199 from django.test import TestCase
200 
201-from .models import Person, Book, Car, PersonManager, PublishedBookManager
202+from .models import Person, Book, Manufacturer, Color, Car, PersonManager, PublishedBookManager
203 
204 
205 class CustomManagerTests(TestCase):
206@@ -42,32 +42,142 @@
207             lambda b: b.title
208         )
209 
210-        c1 = Car.cars.create(name="Corvette", mileage=21, top_speed=180)
211-        c2 = Car.cars.create(name="Neon", mileage=31, top_speed=100)
212+    def test_related_manager(self):
213+        chevy = Manufacturer.objects.create(name="Chevrolet", country="USA")
214+        dodge = Manufacturer.objects.create(name="Dodge", country="USA")
215 
216+        red = Color.objects.create(name="Red", surcharge=False)
217+        white = Color.objects.create(name="White", surcharge=False)
218+        silver = Color.objects.create(name="Silver", surcharge=True)
219+
220+        chevy_corvette = Car.cars.create(name="Corvette", mileage=21, top_speed=180, manufacturer=chevy)
221+        chevy_corvette.available_colors.add(red)
222+        chevy_corvette.available_colors.add(white)
223+
224+        dodge_neon = Car.cars.create(name="Neon", mileage=31, top_speed=100, manufacturer=dodge)
225+        dodge_neon.available_colors.add(white)
226+        dodge_neon.available_colors.add(silver)
227+
228+        dodge_viper = Car.cars.create(name="Viper", mileage=14, top_speed=200, manufacturer=dodge)
229+        dodge_viper.available_colors.add(red)
230+
231         self.assertQuerysetEqual(
232             Car.cars.order_by("name"), [
233                 "Corvette",
234                 "Neon",
235+                "Viper",
236             ],
237             lambda c: c.name
238         )
239 
240         self.assertQuerysetEqual(
241-            Car.fast_cars.all(), [
242+            Car.fast_cars.order_by("name"), [
243                 "Corvette",
244+                "Viper",
245             ],
246             lambda c: c.name
247         )
248 
249+        self.assertQuerysetEqual(
250+            dodge.car_set.order_by("name"), [
251+                "Neon",
252+                "Viper",
253+            ],
254+            lambda c: c.name
255+        )
256+
257+        self.assertQuerysetEqual(
258+            white.car_set.order_by("name"), [
259+                "Corvette",
260+                "Neon",
261+            ],
262+            lambda c: c.name
263+        )
264+
265+        self.assertQuerysetEqual(
266+            dodge_neon.available_colors.order_by("name"), [
267+                "Silver",
268+                "White",
269+            ],
270+            lambda c: c.name
271+        )
272+
273+        # The remove and clear methods must be available. Check only
274+        # for availability here but don't actually run those methods
275+        # -> no "()".
276+        dodge.car_set.remove
277+        dodge.car_set.clear
278+        white.car_set.remove
279+        white.car_set.clear
280+        dodge_neon.available_colors.remove
281+        dodge_neon.available_colors.clear
282+
283         # Each model class gets a "_default_manager" attribute, which is a
284         # reference to the first manager defined in the class. In this case,
285         # it's "cars".
286-
287         self.assertQuerysetEqual(
288             Car._default_manager.order_by("name"), [
289                 "Corvette",
290                 "Neon",
291+                "Viper",
292             ],
293             lambda c: c.name
294         )
295+
296+        # Choosing a custom manager in a reverse relation.
297+        self.assertQuerysetEqual(
298+            dodge.car_set.manager("fast_cars").all(), [
299+                "Viper",
300+            ],
301+            lambda c: c.name
302+        )
303+
304+        # Choosing a custom manager in a many-to-many reverse relation.
305+        self.assertQuerysetEqual(
306+            white.car_set.manager("fast_cars").all(), [
307+                "Corvette",
308+            ],
309+            lambda c: c.name
310+        )
311+
312+        # Choosing a custom manager in many-to-many forward relation.
313+        self.assertQuerysetEqual(
314+            dodge_neon.available_colors.manager("free_of_charge_colors").all(), [
315+                "White",
316+            ],
317+            lambda c: c.name
318+        )
319+
320+        # Unknown managers must raise exception.
321+        self.assertRaises(
322+            AttributeError,
323+            dodge.car_set.manager, "red_cars"
324+        )
325+        self.assertRaises(
326+            AttributeError,
327+            white.car_set.manager, "red_cars"
328+        )
329+        self.assertRaises(
330+            AttributeError,
331+            dodge_neon.available_colors.manager, "beautiful_colors"
332+        )
333+
334+        # Removing from managers with explicitly selected related manager
335+        # must not be possible to avoid mistakes (as the related manager
336+        # does not check if the given objects are actually part of the
337+        # domain selected by that manager).
338+        # clear on (regular) reverse relations is the only exception.
339+        self.assertRaises(AttributeError, lambda: dodge.car_set.manager("fast_cars").remove)
340+        dodge.car_set.manager("fast_cars").clear  # check availability, but don't run method
341+        self.assertRaises(AttributeError, lambda: white.car_set.manager("fast_cars").remove)
342+        self.assertRaises(AttributeError, lambda: white.car_set.manager("fast_cars").clear)
343+        self.assertRaises(AttributeError, lambda: dodge_neon.available_colors.manager("free_of_charge_colors").remove)
344+        self.assertRaises(AttributeError, lambda: dodge_neon.available_colors.manager("free_of_charge_colors").clear)
345+
346+        # Clearing reverse relation with custom manager must only
347+        # remove objects selected by that manager.
348+        self.assertEqual(dodge.car_set.count(), 2)
349+        dodge.car_set.manager("fast_cars").clear()
350+        self.assertEqual(dodge.car_set.count(), 1)
351+        dodge.car_set.clear()
352+        self.assertEqual(dodge.car_set.count(), 0)
353Index: tests/modeltests/custom_managers/models.py
354===================================================================
355--- tests/modeltests/custom_managers/models.py  (revision 17204)
356+++ tests/modeltests/custom_managers/models.py  (working copy)
357@@ -44,12 +44,28 @@
358 
359 # An example of providing multiple custom managers.
360 
361+class Manufacturer(models.Model):
362+    name = models.CharField(max_length=10)
363+    country = models.CharField(max_length=20)
364+
365+class FreeOfChargeColorManager(models.Manager):
366+    def get_query_set(self):
367+        return super(FreeOfChargeColorManager, self).get_query_set().filter(surcharge=False)
368+
369+class Color(models.Model):
370+    name = models.CharField(max_length=10)
371+    surcharge = models.BooleanField()
372+    objects = models.Manager()
373+    free_of_charge_colors = FreeOfChargeColorManager()
374+
375 class FastCarManager(models.Manager):
376     def get_query_set(self):
377         return super(FastCarManager, self).get_query_set().filter(top_speed__gt=150)
378 
379 class Car(models.Model):
380     name = models.CharField(max_length=10)
381+    manufacturer = models.ForeignKey(Manufacturer, null=True)
382+    available_colors = models.ManyToManyField(Color)
383     mileage = models.IntegerField()
384     top_speed = models.IntegerField(help_text="In miles per hour.")
385     cars = models.Manager()