Ticket #4924: compression.2.diff

django/core/management.py

@@ -6 +6 @@
 from optparse import OptionParser
 from django.utils import termcolors
 import os, re, shutil, sys, textwrap
+import bz2, gzip
 
 try:
     set
@@ -1384 +1385 @@
     transaction.enter_transaction_management()
     transaction.managed(True)
 
+    compression_types = {
+        None:   file,
+        'bz2':  bz2.BZ2File,
+        'gz':   gzip.GzipFile,
+    }
+
     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
             formats = serializers.get_serializer_formats()
@@ -1413 +1427 @@
                 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 style.ERROR("Multiple fixtures named '%s' in %s. Aborting." %
-                            (fixture_name, humanize(fixture_dir)))
-                        transaction.rollback()
-                        transaction.leave_transaction_management()
-                        return
-                    else:
-                        count[1] += 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:
-                                count[0] += 1
-                                models.add(obj.object.__class__)
-                                obj.save()
-                            label_found = True
-                        except Exception, e:
+
+                for compression_format in compression_formats:
+                    try:
+                        if compression_format:
+                            full_path = os.path.join(fixture_dir,
+                                '.'.join([fixture_name, format,
+                                compression_format]))
+                        else:
+                            full_path = os.path.join(fixture_dir,
+                                '.'.join([fixture_name, format]))
+
+                        open_method = compression_types[compression_format]
+                        fixture = open_method(full_path, 'r')
+                        if label_found:
                             fixture.close()
-                            sys.stderr.write(
-                                style.ERROR("Problem installing fixture '%s': %s\n" %
-                                     (full_path, str(e))))
+                            print style.ERROR("Multiple fixtures named '%s' in %s. Aborting." %
+                                (fixture_name, humanize(fixture_dir)))
                             transaction.rollback()
                             transaction.leave_transaction_management()
                             return
-                        fixture.close()
-                except:
-                    if verbosity > 1:
-                        print "No %s fixture '%s' in %s." % \
-                            (format, fixture_name, humanize(fixture_dir))
+                        else:
+                            count[1] += 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:
+                                    count[0] += 1
+                                    models.add(obj.object.__class__)
+                                    obj.save()
+                                label_found = True
+                            except Exception, e:
+                                fixture.close()
+                                sys.stderr.write(
+                                    style.ERROR("Problem installing fixture '%s': %s\n" %
+                                         (full_path, str(e))))
+                                transaction.rollback()
+                                transaction.leave_transaction_management()
+                                return
+                            fixture.close()
+                    except:
+                        if verbosity > 1:
+                            print "No %s fixture '%s' in %s." % \
+                                (format, fixture_name, humanize(fixture_dir))
 
     if count[0] > 0:
         sequence_sql = backend.get_sql_sequence_reset(style, models)

tests/modeltests/fixtures/models.py

@@ -59 +59 @@
 >>> Article.objects.all()
 [<Article: Python program becomes self aware>]
 
+# Load fixture 4 (compressed), using format discovery
+>>> management.load_data(['fixture4'], verbosity=0)
+>>> Article.objects.all()
+[<Article: End of the world nigh>, <Article: Relationship trouble? Blame your mother>, <Article: Python program becomes self aware>]
+
 # Load fixture 1 again, using format discovery
+>>> management.flush(verbosity=0, interactive=False)
 >>> management.load_data(['fixture1'], verbosity=0)
 >>> Article.objects.all()
 [<Article: Time to reform copyright>, <Article: Poker has no place on ESPN>, <Article: Python program becomes self aware>]
@@ -69 +75 @@
 >>> management.load_data(['fixture2'], verbosity=0) # doctest: +ELLIPSIS
 Multiple fixtures named 'fixture2' in '...fixtures'. Aborting.
 
+# Try to load fixture 5 using format discovery; this will fail
+# because there are two fixture5's, one compressed with bzip.
+>>> management.load_data(['fixture5'], verbosity=0) # doctest: +ELLIPSIS
+Multiple fixtures named 'fixture5' in '...fixtures'. Aborting.
+
 >>> Article.objects.all()
 [<Article: Time to reform copyright>, <Article: Poker has no place on ESPN>, <Article: Python program becomes self aware>]