Ticket #18020: loaddata_refactor_compression_types.diff

django/core/management/commands/loaddata.py

@@ -18 +18 @@
 from django.db.models import get_apps
 from django.utils.itercompat import product
 
-try:
-    import bz2
-    has_bz2 = True
-except ImportError:
-    has_bz2 = False
 
+def get_module(module_name):
+    try:
+        module = __import__(module_name)
+    except ImportError:
+        return None
+    return module
+
+
 class Command(BaseCommand):
     help = 'Installs the named fixture(s) in the database.'
     args = "fixture [fixture ...]"
@@ -33 +36 @@
             default=DEFAULT_DB_ALIAS, help='Nominates a specific database to load '
                 'fixtures into. Defaults to the "default" database.'),
     )
+ 
+    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])
+   
+    def __init__(self, *args, **kwargs):
+        super(self.__class__, self).__init__(*args, **kwargs)
+        self._compression_types = {
+            None:   open,
+            'gz':   gzip.GzipFile,
+            'zip':  self.__class__._SingleZipReader
+        }
+        bz2 = get_module('bz2')
+        if bz2:
+            self._compression_types['bz2'] = bz2.BZ2File
 
     def handle(self, *fixture_labels, **options):
         using = options.get('database')
@@ -77 +99 @@
             transaction.enter_transaction_management(using=using)
             transaction.managed(True, using=using)
 
-        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:   open,
-            'gz':   gzip.GzipFile,
-            'zip':  SingleZipReader
-        }
-        if has_bz2:
-            compression_types['bz2'] = bz2.BZ2File
-
         app_module_paths = []
         for app in get_apps():
             if hasattr(app, '__path__'):
@@ -110 +116 @@
                 for fixture_label in fixture_labels:
                     parts = fixture_label.split('.')
 
-                    if len(parts) > 1 and parts[-1] in compression_types:
+                    if len(parts) > 1 and parts[-1] in self._compression_types:
                         compression_formats = [parts[-1]]
                         parts = parts[:-1]
                     else:
-                        compression_formats = compression_types.keys()
+                        compression_formats = self._compression_types.keys()
 
                     if len(parts) == 1:
                         fixture_name = parts[0]
@@ -161 +167 @@
                                 self.stdout.write("Trying %s for %s fixture '%s'...\n" % \
                                     (humanize(fixture_dir), file_name, fixture_name))
                             full_path = os.path.join(fixture_dir, file_name)
-                            open_method = compression_types[compression_format]
+                            open_method = self._compression_types[compression_format]
                             try:
                                 fixture = open_method(full_path, 'r')
                             except IOError: