Ticket #3121: get_or_create.diff

django/db/models/fields/related.py

@@ -226 +226 @@
                 return new_obj
             create.alters_data = True
 
+            def get_or_create(self, **kwargs):
+                assert len(kwargs), 'get_or_create() must be passed at least one keyword argument'
+                defaults = kwargs.pop('defaults', {})
+                try:
+                    return self.get(**kwargs), False
+                except self.model.DoesNotExist:
+                    params = dict([(k, v) for k, v in kwargs.items() if '__' not in k])
+                    params.update(defaults)
+                    obj = self.create(**params)
+                    return obj, True
+            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):
@@ -314 +326 @@
             return new_obj
         create.alters_data = True
 
+        def get_or_create(self, **kwargs):
+            assert len(kwargs), 'get_or_create() must be passed at least one keyword argument'
+            defaults = kwargs.pop('defaults', {})
+            try:
+                return self.get(**kwargs), False
+            except self.model.DoesNotExist:
+                params = dict([(k, v) for k, v in kwargs.items() if '__' not in k])
+                params.update(defaults)
+                obj = self.create(**params)
+                return obj, True
+        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/modeltests/get_or_create/models.py

@@ -15 +15 @@
     def __unicode__(self):
         return u'%s %s' % (self.first_name, self.last_name)
 
+class Publisher(models.Model):
+    name = models.CharField(maxlength=100)
+
+class Author(models.Model):
+    name = models.CharField(maxlength=100)
+
+class Book(models.Model):
+    name = models.CharField(maxlength=100)
+    authors = models.ManyToManyField(Author, related_name='books')
+    publisher = models.ForeignKey(Publisher, related_name='books')
+
 __test__ = {'API_TESTS':"""
 # Acting as a divine being, create an Person.
 >>> from datetime import date
@@ -49 +60 @@
 False
 >>> Person.objects.count()
 2
+
+# First create a Publisher.
+>>> p = Publisher.objects.create(name='Acme Publishing')
+
+# Create a book, it should be created.
+>>> book, created = p.books.get_or_create(name='The Book of Ed & Fred')
+>>> created
+True
+
+# Try again.
+>>> book, created = p.books.get_or_create(name='The Book of Ed & Fred')
+>>> created
+False
+
+# A book needs Authors.
+>>> ed, created = book.authors.get_or_create(name='Ed')
+>>> fred, created = book.authors.get_or_create(name='Fred')
+
+# Add another Author, without any books.
+>>> Author.objects.create(name='Ted')
+<Author: Author object>
+
+# By now, there should be 3 Authors in total. The `book' object should have
+# two.
+>>> Author.objects.count()
+3
+>>> book.authors.count()
+2
+
+# Ed writes another Book.
+>>> 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
 """}