From a949f9ec7d8bdd348f628c2d579f4cc633c56e56 Mon Sep 17 00:00:00 2001
From: Gary Wilson Jr <gary.wilson@gmail.com>
Date: Sat, 16 Aug 2008 20:59:06 +0000
Subject: [PATCH] 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            | 17 ++++
 .../get_or_create_regress/__init__.py         |  0
 .../get_or_create_regress/models.py           | 91 +++++++++++++++++++
 3 files changed, 108 insertions(+)
 create mode 100644 tests/regressiontests/get_or_create_regress/__init__.py
 create mode 100644 tests/regressiontests/get_or_create_regress/models.py

diff --git a/django/db/models/fields/related.py b/django/db/models/fields/related.py
index e90c54df4e..e2c78ef96b 100644
--- a/django/db/models/fields/related.py
+++ b/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
diff --git a/tests/regressiontests/get_or_create_regress/__init__.py b/tests/regressiontests/get_or_create_regress/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/regressiontests/get_or_create_regress/models.py b/tests/regressiontests/get_or_create_regress/models.py
new file mode 100644
index 0000000000..1b7b00a69a
--- /dev/null
+++ b/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
+"""}