Ticket #13240: add_remove_querysets_v2.diff

File add_remove_querysets_v2.diff, 7.1 KB (added by gabrielhurley, 5 years ago)

new and improved version that handles iterables (not just querysets)

  • django/db/models/fields/related.py

     
    411411
    412412            def add(self, *objs):
    413413                for obj in objs:
    414                     if not isinstance(obj, self.model):
    415                         raise TypeError("'%s' instance expected" % self.model._meta.object_name)
    416                     setattr(obj, rel_field.name, instance)
    417                     obj.save()
     414                    try:
     415                        for item in iter(obj):
     416                            self.add(item)
     417                    except TypeError:
     418                        if not isinstance(obj, self.model):
     419                            raise TypeError("'%s' instance expected" % self.model._meta.object_name)
     420                        setattr(obj, rel_field.name, instance)
     421                        obj.save()
    418422            add.alters_data = True
    419423
    420424            def create(self, **kwargs):
     
    436440                def remove(self, *objs):
    437441                    val = getattr(instance, rel_field.rel.get_related_field().attname)
    438442                    for obj in objs:
    439                         # Is obj actually part of this descriptor set?
    440                         if getattr(obj, rel_field.attname) == val:
    441                             setattr(obj, rel_field.name, None)
    442                             obj.save()
    443                         else:
    444                             raise rel_field.rel.to.DoesNotExist("%r is not related to %r." % (obj, instance))
     443                        try:
     444                            for item in iter(obj):
     445                                self.remove(item)
     446                        except TypeError:
     447                            # Is obj actually part of this descriptor set?
     448                            if getattr(obj, rel_field.attname) == val:
     449                                setattr(obj, rel_field.name, None)
     450                                obj.save()
     451                            else:
     452                                raise rel_field.rel.to.DoesNotExist("%r is not related to %r." % (obj, instance))
    445453                remove.alters_data = True
    446454
    447455                def clear(self):
     
    551559                        new_ids.add(obj.pk)
    552560                    elif isinstance(obj, Model):
    553561                        raise TypeError("'%s' instance expected" % self.model._meta.object_name)
     562                    elif isinstance(obj, (int, basestring)):
     563                        new_ids.add(obj)
     564                    elif hasattr(obj, "__iter__"):
     565                        self._add_items(source_field_name, target_field_name, *obj)
    554566                    else:
    555                         new_ids.add(obj)
     567                        raise TypeError("%s cannot be added to %s" % (obj, self.model._meta.object_name))
    556568                db = router.db_for_write(self.through.__class__, instance=self.instance)
    557569                vals = self.through._default_manager.using(db).values_list(target_field_name, flat=True)
    558570                vals = vals.filter(**{
     
    592604                for obj in objs:
    593605                    if isinstance(obj, self.model):
    594606                        old_ids.add(obj.pk)
     607                    elif isinstance(obj, (int, basestring)):
     608                        old_ids.add(obj)
     609                    elif hasattr(obj, "__iter__"):
     610                        self._remove_items(source_field_name, target_field_name, *obj)
    595611                    else:
    596                         old_ids.add(obj)
     612                        raise TypeError("%s cannot be added to %s" % (obj, self.model._meta.object_name))
    597613                # Work out what DB we're operating on
    598614                db = router.db_for_write(self.through.__class__, instance=self.instance)
    599615                # Send a signal to the other end if need be.
  • docs/ref/models/relations.txt

     
    3636    .. method:: add(obj1, [obj2, ...])
    3737
    3838        Adds the specified model objects to the related object set.
     39       
     40        .. versionchanged:: 1.3
     41       
     42        The ``add()`` method can now accept any iterable, including``QuerySet`` objects.
    3943
    4044        Example::
    4145
     
    8892        ``e`` from ``b.entry_set()`` is equivalent to doing ``e.blog = None``,
    8993        and because the ``blog`` ``ForeignKey`` doesn't have ``null=True``, this
    9094        is invalid.
     95       
     96        .. versionchanged:: 1.3
     97       
     98        The ``remove()`` method can now accept any iterable, including``QuerySet`` objects.
    9199
    92100    .. method:: clear()
    93101
  • tests/modeltests/many_to_many/models.py

     
    6161# Adding a second time is OK
    6262>>> a2.publications.add(p3)
    6363
     64# Adding and removing using QuerySets and iterables works too
     65>>> p4 = Publication(id=None, title='Science Times')
     66>>> p4.save()
     67>>> science_pubs = Publication.objects.filter(title__contains="Science")
     68>>> a2.publications.add(science_pubs)
     69>>> a2.publications.all()
     70[<Publication: Science News>, <Publication: Science Times>, <Publication: Science Weekly>, <Publication: The Python Journal>]
     71>>> a2.publications.remove(Publication.objects.filter(title__contains="Times"))
     72>>> a2.publications.add(list(science_pubs))
     73>>> a2.publications.all()
     74[<Publication: Science News>, <Publication: Science Times>, <Publication: Science Weekly>, <Publication: The Python Journal>]
     75>>> a2.publications.remove(Publication.objects.filter(title__contains="Times"))
     76>>> p4.delete()
     77
    6478# Adding an object of the wrong type raises TypeError
    6579>>> a2.publications.add(a1)
    6680Traceback (most recent call last):
  • tests/modeltests/many_to_one/models.py

     
    7979...
    8080TypeError: 'Article' instance expected
    8181
     82>>> r3 = Reporter(first_name='Eric', last_name='James', email='eric@example.com')
     83>>> r3.save()
     84>>> a3 = Article.objects.create(headline="Human interest story", reporter=r, pub_date=datetime(2010, 3, 04))
     85>>> a4 = Article.objects.create(headline="Eric's story", reporter=r, pub_date=datetime(2010, 3, 14))
     86>>> a5 = Article.objects.create(headline="Wild Parties!", reporter=r, pub_date=datetime(2010, 9, 22))
     87>>> new_articles = Article.objects.filter(pub_date__year=2010, pub_date__month=3)
     88>>> r3.article_set.add(new_articles)
     89>>> r3.article_set.all()
     90[<Article: Eric's story>, <Article: Human interest story>]
     91>>> r3.article_set.add([a5])
     92>>> r3.article_set.all()
     93[<Article: Eric's story>, <Article: Human interest story>, <Article: Wild Parties!>]
     94>>> new_articles.delete()
     95>>> a5.delete()
     96>>> r3.delete()
     97
    8298>>> r.article_set.all()
    8399[<Article: John's second story>, <Article: This is a test>]
    84100>>> r2.article_set.all()
Back to Top