Refactor loaddata for readability. (9e7723f8) · Commits · Dom Sekotill / django

django/core/management/commands/loaddata.py

+140 −133

Original line number	Diff line number	Diff line
		from __future__ import unicode_literals

		import sys
		import os
		import gzip
		import zipfile
		from optparse import make_option
		import traceback

		from django.conf import settings
		from django.core import serializers
		@@ -39,10 +37,10 @@ class Command(BaseCommand):

		def handle(self, fixture_labels, *options):

		ignore = options.get('ignore')
		using = options.get('database')
		self.ignore = options.get('ignore')
		self.using = options.get('database')

		connection = connections[using]
		connection = connections[self.using]

		if not len(fixture_labels):
		raise CommandError(
		@@ -50,8 +48,7 @@ class Command(BaseCommand):
		"least one fixture in the command line."
		)

		verbosity = int(options.get('verbosity'))
		show_traceback = options.get('traceback')
		self.verbosity = int(options.get('verbosity'))

		# commit is a stealth option - it isn't really useful as
		# a command line option, but it can be useful when invoking
		@@ -62,12 +59,10 @@ class Command(BaseCommand):
		commit = options.get('commit', True)

		# Keep a count of the installed objects and fixtures
		fixture_count = 0
		loaded_object_count = 0
		fixture_object_count = 0
		models = set()

		humanize = lambda dirname: "'%s'" % dirname if dirname else 'absolute path'
		self.fixture_count = 0
		self.loaded_object_count = 0
		self.fixture_object_count = 0
		self.models = set()

		# Get a cursor (even though we don't need one yet). This has
		# the side effect of initializing the test database (if
		@@ -77,9 +72,9 @@ class Command(BaseCommand):
		# Start transaction management. All fixtures are installed in a
		# single transaction to ensure that all references are resolved.
		if commit:
		transaction.commit_unless_managed(using=using)
		transaction.enter_transaction_management(using=using)
		transaction.managed(True, using=using)
		transaction.commit_unless_managed(using=self.using)
		transaction.enter_transaction_management(using=self.using)
		transaction.managed(True, using=self.using)

		class SingleZipReader(zipfile.ZipFile):
		def __init__(self, args, *kwargs):
		@@ -89,13 +84,13 @@ class Command(BaseCommand):
		def read(self):
		return zipfile.ZipFile.read(self, self.namelist()[0])

		compression_types = {
		self.compression_types = {
		None: open,
		'gz': gzip.GzipFile,
		'zip': SingleZipReader
		}
		if has_bz2:
		compression_types['bz2'] = bz2.BZ2File
		self.compression_types['bz2'] = bz2.BZ2File

		app_module_paths = []
		for app in get_apps():
		@@ -112,13 +107,63 @@ class Command(BaseCommand):
		try:
		with connection.constraint_checks_disabled():
		for fixture_label in fixture_labels:
		self.load_label(fixture_label, app_fixtures)

		# Since we disabled constraint checks, we must manually check for
		# any invalid keys that might have been added
		table_names = [model._meta.db_table for model in self.models]
		try:
		connection.check_constraints(table_names=table_names)
		except Exception as e:
		e.args = ("Problem installing fixtures: %s" % e,)
		raise

		except (SystemExit, KeyboardInterrupt):
		raise
		except Exception as e:
		if commit:
		transaction.rollback(using=self.using)
		transaction.leave_transaction_management(using=self.using)
		raise

		# If we found even one object in a fixture, we need to reset the
		# database sequences.
		if self.loaded_object_count > 0:
		sequence_sql = connection.ops.sequence_reset_sql(no_style(), self.models)
		if sequence_sql:
		if self.verbosity >= 2:
		self.stdout.write("Resetting sequences\n")
		for line in sequence_sql:
		cursor.execute(line)

		if commit:
		transaction.commit(using=self.using)
		transaction.leave_transaction_management(using=self.using)

		if self.verbosity >= 1:
		if self.fixture_object_count == self.loaded_object_count:
		self.stdout.write("Installed %d object(s) from %d fixture(s)" % (
		self.loaded_object_count, self.fixture_count))
		else:
		self.stdout.write("Installed %d object(s) (of %d) from %d fixture(s)" % (
		self.loaded_object_count, self.fixture_object_count, self.fixture_count))

		# Close the DB connection. This is required as a workaround for an
		# edge case in MySQL: if the same connection is used to
		# create tables, load data, and query, the query can return
		# incorrect results. See Django #7572, MySQL #37735.
		if commit:
		connection.close()

		def load_label(self, fixture_label, app_fixtures):

		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]
		@@ -131,7 +176,7 @@ class Command(BaseCommand):
		formats = []

		if formats:
		if verbosity >= 2:
		if self.verbosity >= 2:
		self.stdout.write("Loading '%s' fixtures..." % fixture_name)
		else:
		raise CommandError(
		@@ -144,11 +189,19 @@ class Command(BaseCommand):
		fixture_dirs = app_fixtures + list(settings.FIXTURE_DIRS) + ['']

		for fixture_dir in fixture_dirs:
		if verbosity >= 2:
		self.process_dir(fixture_dir, fixture_name, compression_formats,
		formats)

		def process_dir(self, fixture_dir, fixture_name, compression_formats,
		serialization_formats):

		humanize = lambda dirname: "'%s'" % dirname if dirname else 'absolute path'

		if self.verbosity >= 2:
		self.stdout.write("Checking %s for fixtures..." % humanize(fixture_dir))

		label_found = False
		for combo in product([using, None], formats, compression_formats):
		for combo in product([self.using, None], serialization_formats, compression_formats):
		database, format, compression_format = combo
		file_name = '.'.join(
		p for p in [
		@@ -157,15 +210,15 @@ class Command(BaseCommand):
		if p
		)

		if verbosity >= 3:
		if self.verbosity >= 3:
		self.stdout.write("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]
		open_method = self.compression_types[compression_format]
		try:
		fixture = open_method(full_path, 'r')
		except IOError:
		if verbosity >= 2:
		if self.verbosity >= 2:
		self.stdout.write("No %s fixture '%s' in %s." % \
		(format, fixture_name, humanize(fixture_dir)))
		else:
		@@ -174,22 +227,22 @@ class Command(BaseCommand):
		raise CommandError("Multiple fixtures named '%s' in %s. Aborting." %
		(fixture_name, humanize(fixture_dir)))

		fixture_count += 1
		self.fixture_count += 1
		objects_in_fixture = 0
		loaded_objects_in_fixture = 0
		if verbosity >= 2:
		if self.verbosity >= 2:
		self.stdout.write("Installing %s fixture '%s' from %s." % \
		(format, fixture_name, humanize(fixture_dir)))

		objects = serializers.deserialize(format, fixture, using=using, ignorenonexistent=ignore)
		objects = serializers.deserialize(format, fixture, using=self.using, ignorenonexistent=self.ignore)

		for obj in objects:
		objects_in_fixture += 1
		if router.allow_syncdb(using, obj.object.__class__):
		if router.allow_syncdb(self.using, obj.object.__class__):
		loaded_objects_in_fixture += 1
		models.add(obj.object.__class__)
		self.models.add(obj.object.__class__)
		try:
		obj.save(using=using)
		obj.save(using=self.using)
		except (DatabaseError, IntegrityError) as e:
		e.args = ("Could not load %(app_label)s.%(object_name)s(pk=%(pk)s): %(error_msg)s" % {
		'app_label': obj.object._meta.app_label,
		@@ -199,8 +252,8 @@ class Command(BaseCommand):
		},)
		raise

		loaded_object_count += loaded_objects_in_fixture
		fixture_object_count += objects_in_fixture
		self.loaded_object_count += loaded_objects_in_fixture
		self.fixture_object_count += objects_in_fixture
		label_found = True
		except Exception as e:
		if not isinstance(e, CommandError):
		@@ -215,49 +268,3 @@ class Command(BaseCommand):
		raise CommandError(
		"No fixture data found for '%s'. (File format may be invalid.)" %
		(fixture_name))

		# Since we disabled constraint checks, we must manually check for
		# any invalid keys that might have been added
		table_names = [model._meta.db_table for model in models]
		try:
		connection.check_constraints(table_names=table_names)
		except Exception as e:
		e.args = ("Problem installing fixtures: %s" % e,)
		raise

		except (SystemExit, KeyboardInterrupt):
		raise
		except Exception as e:
		if commit:
		transaction.rollback(using=using)
		transaction.leave_transaction_management(using=using)
		raise

		# If we found even one object in a fixture, we need to reset the
		# database sequences.
		if loaded_object_count > 0:
		sequence_sql = connection.ops.sequence_reset_sql(no_style(), models)
		if sequence_sql:
		if verbosity >= 2:
		self.stdout.write("Resetting sequences\n")
		for line in sequence_sql:
		cursor.execute(line)

		if commit:
		transaction.commit(using=using)
		transaction.leave_transaction_management(using=using)

		if verbosity >= 1:
		if fixture_object_count == loaded_object_count:
		self.stdout.write("Installed %d object(s) from %d fixture(s)" % (
		loaded_object_count, fixture_count))
		else:
		self.stdout.write("Installed %d object(s) (of %d) from %d fixture(s)" % (
		loaded_object_count, fixture_object_count, fixture_count))

		# Close the DB connection. This is required as a workaround for an
		# edge case in MySQL: if the same connection is used to
		# create tables, load data, and query, the query can return
		# incorrect results. See Django #7572, MySQL #37735.
		if commit:
		connection.close()