Fixed #4924: added support for loading compressed fixtures. Thanks to Lars…

Fixed #4924: added support for loading compressed fixtures. Thanks to Lars Yencken and Jeremy Dunck.

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

AUTHORS

@@ -425,6 +425,7 @@ answer newbie questions, and generally made Django that much better:
     Maciej Wiśniowski <pigletto@gmail.com>
     wojtek
     Jason Yan <tailofthesun@gmail.com>
+    Lars Yencken <lars.yencken@gmail.com>
     ye7cakf02@sneakemail.com
     ymasuda@ethercube.com
     Jesse Young <adunar@gmail.com>

django/core/management/commands/loaddata.py

@@ -3,6 +3,7 @@ from django.core.management.color import no_style
 from optparse import make_option
 import sys
 import os
+import bz2, gzip, zipfile
 
 try:
     set
@@ -51,11 +52,33 @@ class Command(BaseCommand):
             transaction.enter_transaction_management()
             transaction.managed(True)
 
+        class SingleZipReader(zipfile.ZipFile):
+            def __init__(self, *args, **kwargs):
+                zipfile.ZipFile.__init__(self, *args, **kwargs)
+                if settings.DEBUG:
+                    assert len(self.namelist()) == 1, "Zip-compressed fixtures must contain only one file."
+            def read(self):
+                return zipfile.ZipFile.read(self, self.namelist()[0])
+
+        compression_types = {
+            None:   file,
+            'bz2':  bz2.BZ2File,
+            'gz':   gzip.GzipFile,
+            'zip':  SingleZipReader
+        }
+
         app_fixtures = [os.path.join(os.path.dirname(app.__file__), 'fixtures') for app in get_apps()]
         for fixture_label in fixture_labels:
             parts = fixture_label.split('.')
+
+            if len(parts) > 1 and parts[-1] in compression_types:
+                compression_formats = [parts[-1]]
+                parts = parts[:-1]
+            else:
+                compression_formats = compression_types.keys()
+
             if len(parts) == 1:
-                fixture_name = fixture_label
+                fixture_name = parts[0]
                 formats = serializers.get_public_serializer_formats()
             else:
                 fixture_name, format = '.'.join(parts[:-1]), parts[-1]
@@ -86,64 +109,63 @@ class Command(BaseCommand):
 
                 label_found = False
                 for format in formats:
-                    serializer = serializers.get_serializer(format)
-                    if verbosity > 1:
-                        print "Trying %s for %s fixture '%s'..." % \
-                            (humanize(fixture_dir), format, fixture_name)
-                    try:
-                        full_path = os.path.join(fixture_dir, '.'.join([fixture_name, format]))
-                        fixture = open(full_path, 'r')
-                        if label_found:
-                            fixture.close()
-                            print self.style.ERROR("Multiple fixtures named '%s' in %s. Aborting." %
-                                (fixture_name, humanize(fixture_dir)))
-                            transaction.rollback()
-                            transaction.leave_transaction_management()
-                            return
-                        else:
-                            fixture_count += 1
-                            objects_in_fixture = 0
-                            if verbosity > 0:
-                                print "Installing %s fixture '%s' from %s." % \
-                                    (format, fixture_name, humanize(fixture_dir))
-                            try:
-                                objects = serializers.deserialize(format, fixture)
-                                for obj in objects:
-                                    objects_in_fixture += 1
-                                    models.add(obj.object.__class__)
-                                    obj.save()
-                                object_count += objects_in_fixture
-                                label_found = True
-                            except (SystemExit, KeyboardInterrupt):
-                                raise
-                            except Exception:
-                                import traceback
+                    for compression_format in compression_formats:
+                        if compression_format: 
+                            file_name = '.'.join([fixture_name, format, 
+                                                  compression_format])
+                        else: 
+                            file_name = '.'.join([fixture_name, format])
+                    
+                        if verbosity > 1:
+                            print "Trying %s for %s fixture '%s'..." % \
+                                (humanize(fixture_dir), file_name, fixture_name)
+                        full_path = os.path.join(fixture_dir, file_name)
+                        open_method = compression_types[compression_format]                                
+                        try: 
+                            fixture = open_method(full_path, 'r')
+                            if label_found:
                                 fixture.close()
+                                print self.style.ERROR("Multiple fixtures named '%s' in %s. Aborting." %
+                                    (fixture_name, humanize(fixture_dir)))
                                 transaction.rollback()
                                 transaction.leave_transaction_management()
-                                if show_traceback:
-                                    import traceback
-                                    traceback.print_exc()
-                                else:
-                                    sys.stderr.write(
-                                        self.style.ERROR("Problem installing fixture '%s': %s\n" %
-                                             (full_path, traceback.format_exc())))
                                 return
-                            fixture.close()
-
-                            # If the fixture we loaded contains 0 objects, assume that an
-                            # error was encountered during fixture loading.
-                            if objects_in_fixture == 0:
-                                sys.stderr.write(
-                                    self.style.ERROR("No fixture data found for '%s'. (File format may be invalid.)" %
-                                        (fixture_name)))
+                            else:
+                                fixture_count += 1
+                                if verbosity > 0:
+                                    print "Installing %s fixture '%s' from %s." % \
+                                        (format, fixture_name, humanize(fixture_dir))
+                                try:
+                                    objects = serializers.deserialize(format, fixture)
+                                    for obj in objects:
+                                        object_count += 1
+                                        models.add(obj.object.__class__)
+                                        obj.save()
+                                    label_found = True
+                                except (SystemExit, KeyboardInterrupt):
+                                    raise
+                                except Exception:
+                                    import traceback
+                                    fixture.close()
+                                    transaction.rollback()
+                                    transaction.leave_transaction_management()
+                                    if show_traceback:
+                                        import traceback
+                                        traceback.print_exc()
+                                    else:
+                                        sys.stderr.write(
+                                            self.style.ERROR("Problem installing fixture '%s': %s\n" %
+                                                 (full_path, traceback.format_exc())))
+                                    return
+                                fixture.close()
+                        except Exception, e:
+                            if verbosity > 1:
+                                print "No %s fixture '%s' in %s." % \
+                                    (format, fixture_name, humanize(fixture_dir))
+                                print e
                                 transaction.rollback()
                                 transaction.leave_transaction_management()
                                 return
-                    except:
-                        if verbosity > 1:
-                            print "No %s fixture '%s' in %s." % \
-                                (format, fixture_name, humanize(fixture_dir))
 
         # If we found even one object in a fixture, we need to reset the
         # database sequences.

docs/ref/django-admin.txt

@@ -290,6 +290,9 @@ loaddata <fixture fixture ...>
 
 Searches for and loads the contents of the named fixture into the database.
 
+What's a "fixture"?
+~~~~~~~~~~~~~~~~~~~
+
 A *fixture* is a collection of files that contain the serialized contents of
 the database. Each fixture has a unique name, and the files that comprise the
 fixture can be distributed over multiple directories, in multiple applications.
@@ -309,21 +312,17 @@ will be loaded. For example::
     django-admin.py loaddata mydata.json
 
 would only load JSON fixtures called ``mydata``. The fixture extension
-must correspond to the registered name of a serializer (e.g., ``json`` or
-``xml``).
+must correspond to the registered name of a
+:ref:`serializer <serialization-formats>` (e.g., ``json`` or ``xml``).
 
-If you omit the extension, Django will search all available fixture types
+If you omit the extensions, Django will search all available fixture types
 for a matching fixture. For example::
 
     django-admin.py loaddata mydata
 
 would look for any fixture of any fixture type called ``mydata``. If a fixture
 directory contained ``mydata.json``, that fixture would be loaded
-as a JSON fixture. However, if two fixtures with the same name but different
-fixture type are discovered (for example, if ``mydata.json`` and
-``mydata.xml`` were found in the same fixture directory), fixture
-installation will be aborted, and any data installed in the call to
-``loaddata`` will be removed from the database.
+as a JSON fixture.
 
 The fixtures that are named can include directory components. These
 directories will be included in the search path. For example::
@@ -342,6 +341,23 @@ end of the transaction.
 
 The ``dumpdata`` command can be used to generate input for ``loaddata``.
 
+Compressed fixtures
+~~~~~~~~~~~~~~~~~~~
+
+Fixtures may be compressed in ``zip``, ``gz``, or ``bz2`` format. For example::
+
+    django-admin.py loaddata mydata.json
+
+would look for any of ``mydata.json``, ``mydata.json.zip``,
+``mydata.json.gz``, or ``mydata.json.bz2``.  The first file contained within a
+zip-compressed archive is used.
+
+Note that if two fixtures with the same name but different
+fixture type are discovered (for example, if ``mydata.json`` and
+``mydata.xml.gz`` were found in the same fixture directory), fixture
+installation will be aborted, and any data installed in the call to
+``loaddata`` will be removed from the database.
+
 .. admonition:: MySQL and Fixtures
 
     Unfortunately, MySQL isn't capable of completely supporting all the

tests/modeltests/fixtures/models.py

@@ -76,12 +76,46 @@ if settings.DATABASE_ENGINE != 'mysql':
 >>> management.call_command('loaddata', 'fixture2', verbosity=0) # doctest: +ELLIPSIS
 Multiple fixtures named 'fixture2' in '...fixtures'. Aborting.
 
+# object list is unaffected
 >>> Article.objects.all()
 [<Article: Time to reform copyright>, <Article: Poker has no place on ESPN>, <Article: Python program becomes self aware>]
 
 # Dump the current contents of the database as a JSON fixture
 >>> management.call_command('dumpdata', 'fixtures', format='json')
 [{"pk": 3, "model": "fixtures.article", "fields": {"headline": "Time to reform copyright", "pub_date": "2006-06-16 13:00:00"}}, {"pk": 2, "model": "fixtures.article", "fields": {"headline": "Poker has no place on ESPN", "pub_date": "2006-06-16 12:00:00"}}, {"pk": 1, "model": "fixtures.article", "fields": {"headline": "Python program becomes self aware", "pub_date": "2006-06-16 11:00:00"}}]
+
+# Load fixture 4 (compressed), using format discovery
+>>> management.call_command('loaddata', 'fixture4', verbosity=0)
+>>> Article.objects.all()
+[<Article: Django pets kitten>, <Article: Time to reform copyright>, <Article: Poker has no place on ESPN>, <Article: Python program becomes self aware>]
+
+>>> management.call_command('flush', verbosity=0, interactive=False)
+
+# Load fixture 4 (compressed), using format specification
+>>> management.call_command('loaddata', 'fixture4.json', verbosity=0)
+>>> Article.objects.all()
+[<Article: Django pets kitten>, <Article: Python program becomes self aware>]
+
+>>> management.call_command('flush', verbosity=0, interactive=False)
+
+# Load fixture 5 (compressed), using format *and* compression specification
+>>> management.call_command('loaddata', 'fixture5.json.zip', verbosity=0)
+>>> Article.objects.all()
+[<Article: WoW subscribers now outnumber readers>, <Article: Python program becomes self aware>]
+
+>>> management.call_command('flush', verbosity=0, interactive=False)
+
+# Load fixture 5 (compressed), only compression specification
+>>> management.call_command('loaddata', 'fixture5.zip', verbosity=0)
+>>> Article.objects.all()
+[<Article: WoW subscribers now outnumber readers>, <Article: Python program becomes self aware>]
+
+>>> management.call_command('flush', verbosity=0, interactive=False)
+
+# Try to load fixture 5 using format and compression discovery; this will fail
+# because there are two fixture5's in the fixtures directory
+>>> management.call_command('loaddata', 'fixture5', verbosity=0) # doctest: +ELLIPSIS
+Multiple fixtures named 'fixture5' in '...fixtures'. Aborting.
 """
 
 from django.test import TestCase