Fixed #3121 -- Made `get_or_create()` work for `RelatedManager` and `ManyRelatedManager`.

git-svn-id: http://code.djangoproject.com/svn/django/trunk@8415 bcc190cf-cafb-0310-a4f2-bffc1f526a37

django/db/models/fields/related.py

@@ -311,6 +311,13 @@ class ForeignRelatedObjectsDescriptor(object):
                 return new_obj
             create.alters_data = True
 
+            def get_or_create(self, **kwargs):
+                # Update kwargs with the related object that this
+                # ForeignRelatedObjectsDescriptor knows about.
+                kwargs.update({rel_field.name: instance})
+                return super(RelatedManager, self).get_or_create(**kwargs)
+            get_or_create.alters_data = True
+
             # remove() and clear() are only provided if the ForeignKey can have a value of null.
             if rel_field.null:
                 def remove(self, *objs):
@@ -409,6 +416,16 @@ def create_many_related_manager(superclass, through=False):
             return new_obj
         create.alters_data = True
 
+        def get_or_create(self, **kwargs):
+            obj, created = \
+                    super(ManyRelatedManager, self).get_or_create(**kwargs)
+            # We only need to add() if created because if we got an object back
+            # from get() then the relationship already exists.
+            if created:
+                self.add(obj)
+            return obj, created
+        get_or_create.alters_data = True
+
         def _add_items(self, source_col_name, target_col_name, *objs):
             # join_table: name of the m2m link table
             # source_col_name: the PK colname in join_table for the source object

tests/regressiontests/get_or_create_regress/models.py

@@ -0,0 +1,91 @@
+from django.db import models
+
+class Publisher(models.Model):
+    name = models.CharField(max_length=100)
+
+class Author(models.Model):
+    name = models.CharField(max_length=100)
+
+class Book(models.Model):
+    name = models.CharField(max_length=100)
+    authors = models.ManyToManyField(Author, related_name='books')
+    publisher = models.ForeignKey(Publisher, related_name='books')
+
+
+__test__ = {'one':"""
+#
+# RelatedManager
+#
+
+# First create a Publisher.
+>>> p = Publisher.objects.create(name='Acme Publishing')
+
+# Create a book through the publisher.
+>>> book, created = p.books.get_or_create(name='The Book of Ed & Fred')
+>>> created
+True
+
+# The publisher should have one book.
+>>> p.books.count()
+1
+
+# Try get_or_create again, this time nothing should be created.
+>>> book, created = p.books.get_or_create(name='The Book of Ed & Fred')
+>>> created
+False
+
+# And the publisher should still have one book.
+>>> p.books.count()
+1
+
+#
+# ManyRelatedManager
+#
+
+# Add an author to the book.
+>>> ed, created = book.authors.get_or_create(name='Ed')
+>>> created
+True
+
+# Book should have one author.
+>>> book.authors.count()
+1
+
+# Try get_or_create again, this time nothing should be created.
+>>> ed, created = book.authors.get_or_create(name='Ed')
+>>> created
+False
+
+# And the book should still have one author.
+>>> book.authors.count()
+1
+
+# Add a second author to the book.
+>>> fred, created = book.authors.get_or_create(name='Fred')
+>>> created
+True
+
+# The book should have two authors now.
+>>> book.authors.count()
+2
+
+# Create an Author not tied to any books.
+>>> Author.objects.create(name='Ted')
+<Author: Author object>
+
+# There should be three Authors in total. The book object should have two.
+>>> Author.objects.count()
+3
+>>> book.authors.count()
+2
+
+# Try creating a book through an author.
+>>> ed.books.get_or_create(name="Ed's Recipies", publisher=p)
+(<Book: Book object>, True)
+
+# Now Ed has two Books, Fred just one.
+>>> ed.books.count()
+2
+>>> fred.books.count()
+1
+"""}