Ticket #4027: django-model-copying.diff

django/db/models/base.py

@@ -11 +11 @@
 import django.db.models.manager     # Imported to register signal handler.
 from django.core.exceptions import ObjectDoesNotExist, MultipleObjectsReturned, FieldError
 from django.db.models.fields import AutoField, FieldDoesNotExist
-from django.db.models.fields.related import OneToOneRel, ManyToOneRel, OneToOneField
+from django.db.models.fields.related import OneToOneRel, ManyToOneRel, OneToOneField, ForeignKey
 from django.db.models.query import delete_objects, Q
 from django.db.models.query_utils import CollectedObjects, DeferredAttribute
 from django.db.models.options import Options
@@ -606 +606 @@
     def prepare_database_save(self, unused):
         return self.pk
 
+    def _copy(self, recurse, recurse_m2m, copied_objects):
+        if recurse_m2m is None:
+            recurse_m2m = recurse
+
+        # Handle non-m2m fields.
+        self_referencing_fields = []
+        foreign_keys_to_recurse = {}
+        values = {}
+        for field in self._meta.fields:
+            if field.name == self._meta.pk.name:
+                continue
+
+            value = getattr(self, field.name)
+
+            if isinstance(field, ForeignKey) and (value is not None):
+                if value == self:
+                    # This is a reference to the self instance.  Make a note of
+                    # it so that we can correctly refer to the copy of self
+                    # that we're making.  We temporarily store a reference to
+                    # the old self instance to avoid issues where the foreign
+                    # key is not allowed to be NULL.
+                    values[field.name ] = self
+                    self_referencing_fields.append(field)
+                    continue
+
+                if value in copied_objects:
+                    # This object has already been copied.  Store a reference
+                    # to the copy.
+                    values[field.name] = copied_objects[value]
+                    continue
+
+                if recurse:
+                    # We are recursing and this object has not been previously
+                    # copied.  We will make a new copy once the copy of self
+                    # has been created (we need to add that to copied_objects
+                    # before we start recursing).  For now, we store a
+                    # reference to the old value (to avoid issues with NOT
+                    # NULL) and make a note that we need to finish handling
+                    # this field.
+                    values[field.name] = value
+                    foreign_keys_to_recurse[field] = value
+                    continue
+
+            # This is not a foreign key (or it is but its value is None).
+            values[field.name] = value
+
+        copied_self = self.__class__.objects.create(**values)
+        copied_objects[self] = copied_self
+
+        # Now that we have an id for the copy we are making, fill in the values
+        # for self-referencing foreign keys.
+        if self_referencing_fields:
+            for field in self_referencing_fields:
+                setattr(copied_self, field.name, copied_self)
+
+        # self is now in copied_objects, so we can freely recurse on foreign
+        # keys.
+        if foreign_keys_to_recurse:
+            for field, value in foreign_keys_to_recurse.items():
+                copied_value = value._copy(
+                  recurse, recurse_m2m, copied_objects)
+                setattr(copied_self, field.name, copied_value)
+
+        if self_referencing_fields or foreign_keys_to_recurse:
+            copied_self.save()
+
+        # Handle many-to-many relationships.
+        for m2m_field in self._meta.many_to_many:
+            value = getattr(self, m2m_field.attname).all()
+
+            new_value = []
+            for item in value:
+                if item == self:
+                    # This is a link back to the self instance.  Store a
+                    # reference to the copy.
+                    new_value.append(copied_self)
+                    continue
+
+                if (m2m_field.rel.to is self.__class__) and (
+                  item in copied_objects):
+                    # This m2m relationship is to self and we've already begun
+                    # copying the item refered to.  No need to copy the
+                    # reference to that item, as the relationship only needs to
+                    # be specified from one side.
+                    continue
+
+                if item in copied_objects:
+                    # This item has previously been copied.  Store a reference
+                    # to the copy.
+                    new_value.append(copied_objects[item])
+                    continue
+
+                if recurse_m2m:
+                    # This item has not previously been copied and we're
+                    # recursing.  We need a new copy of this item.
+                    new_value.append(
+                      item._copy(recurse, recurse_m2m, copied_objects))
+                else:
+                    # This item has not previously been copied and we're not
+                    # recursing.  Copy the m2m reference to the same item.
+                    new_value.append(item)
+
+            setattr(copied_self, m2m_field.attname, new_value)
+
+        return copied_self
+
+    def copy(self, recurse = False, recurse_m2m = None):
+        '''
+        Creates a copy of this object in the database.
+
+        If recurse is True, objects referenced by foreign keys and many-to-many
+        fields will be recursively copied.  For each object referenced, exactly
+        one copy will be made.  References will be updated such that references
+        to objects that were copied bill be updated to point to the copies.
+
+        If recurse is True but recurse_m2m is False, only objects referenced by
+        foreign keys will be recursed into (objects referenced by many-to-many
+        fields will be left uncopied).  If recurse_m2m is True but recurse is
+        False, only many-to-many fields will be recursed into.
+
+        recurse defaults to False.  recurse_m2m defaults to the value of
+        recurse.  Typically, recurse_m2m will be left unspecified.
+        '''
+        return self._copy(recurse, recurse_m2m, {})
+
 
 ############################################
 # HELPER FUNCTIONS (CURRIED MODEL METHODS) #

tests/modeltests/model_copying/models.py

@@ +1 @@
+try:
+    set
+except NameError:
+    from sets import Set as set
+
+from django.db import models
+
+
+class User(models.Model):
+    username = models.CharField(max_length = 20)
+
+    def __unicode__(self):
+        return self.username
+
+
+class Issue(models.Model):
+    summary = models.CharField(max_length = 256)
+    cc = models.ManyToManyField(
+      User,
+      blank = True,
+      related_name = 'test_issue_cc',
+    )
+    client = models.ForeignKey(User, related_name = 'test_issue_client')
+    duplicate_of = models.ForeignKey('self', blank = True, null = True)
+    related_issues = models.ManyToManyField('self', blank = True)
+
+    def __unicode__(self):
+        return u'#%u: %s' % (self.id, self.summary)
+
+    class Meta:
+        ordering = ('id',)
+
+
+__test__ = {'API_TESTS': """
+>>> foo = User.objects.create(username = 'foo')
+>>> bar = User.objects.create(username = 'bar')
+>>> baz = User.objects.create(username = 'baz')
+
+>>> one = Issue.objects.create(summary = 'A wild and crazy issue', client = foo)
+>>> one.cc = [bar, baz]
+>>> one.id
+1L
+
+# Test non-recursive copying:
+>>> two = one.copy()
+>>> two
+<Issue: #2: A wild and crazy issue>
+>>> two.client == foo
+True
+>>> two.summary == one.summary
+True
+>>> list(two.cc.all()) == list(one.cc.all())
+True
+
+>>> two.delete()
+>>> Issue.objects.all()
+[<Issue: #1: A wild and crazy issue>]
+>>> User.objects.all()
+[<User: foo>, <User: bar>, <User: baz>]
+
+# Test full recursion:
+>>> two = one.copy(recurse = True)
+>>> two
+<Issue: #2: A wild and crazy issue>
+
+# two.client is a copy of foo; it is unequal because it has a different pk
+>>> two.client
+<User: foo>
+>>> two.client != foo
+True
+
+# two.cc is a set of copies of one.cc; they are unequal
+>>> two.cc.all()
+[<User: bar>, <User: baz>]
+>>> set(one.cc.all()) != set(two.cc.all())
+True
+
+>>> two.cc.all().delete()
+>>> two.client.delete()
+>>> two.delete()
+>>> Issue.objects.all()
+[<Issue: #1: A wild and crazy issue>]
+>>> User.objects.all()
+[<User: foo>, <User: bar>, <User: baz>]
+
+# Test recursion with foreign keys only:
+>>> two = one.copy(recurse = True, recurse_m2m = False)
+>>> two
+<Issue: #2: A wild and crazy issue>
+
+# two.client is a copy of foo; it is unequal because it has a different pk
+>>> two.client
+<User: foo>
+>>> two.client != foo
+True
+
+# two.cc is equal to one.cc, because m2m fields were not recursed into
+>>> set(one.cc.all()) == set(two.cc.all())
+True
+
+>>> two.client.delete()
+>>> two.delete()
+>>> Issue.objects.all()
+[<Issue: #1: A wild and crazy issue>]
+>>> User.objects.all()
+[<User: foo>, <User: bar>, <User: baz>]
+
+# Test foreign key recursion with reference to self (this is bogus data and it
+# doesn't make much sense, but it does test the scenario correctly):
+>>> one.duplicate_of = one
+>>> one.save()
+>>> two = one.copy(recurse = True, recurse_m2m = False)
+
+# two.duplicate_of should refer to itself, *not* to one:
+>>> two.duplicate_of
+<Issue: #2: A wild and crazy issue>
+
+>>> two.client.delete()
+>>> two.delete()
+>>> Issue.objects.all()
+[<Issue: #1: A wild and crazy issue>]
+>>> User.objects.all()
+[<User: foo>, <User: bar>, <User: baz>]
+
+# Restore one to its original state:
+>>> one.duplicate_of = None
+>>> one.save()
+
+# Test m2m recursion with reference to self (again, pretty bogus data):
+>>> one.related_issues = [one]
+
+>>> two = one.copy(recurse = False, recurse_m2m = True)
+
+# Again, two should refer to itself, *not* to one:
+>>> two.related_issues.all()
+[<Issue: #2: A wild and crazy issue>]
+"""}