From c296e55dc6a697c7e4d4be92b954633cda4a79b1 Mon Sep 17 00:00:00 2001 From: Claude Paroz Date: Tue, 21 Jul 2015 23:24:32 +0200 Subject: [PATCH] Fixed #22258 -- Added progress status for dumpdata when outputting to file Thanks Gwildor Sok for the report and Tim Graham for the review. --- django/core/management/base.py | 5 +++- django/core/management/commands/dumpdata.py | 23 ++++++++++++---- django/core/serializers/base.py | 30 ++++++++++++++++++++- docs/ref/django-admin.txt | 6 +++++ docs/releases/1.9.txt | 2 ++ tests/fixtures/tests.py | 26 ++++++++++++++++++ tests/serializers/tests.py | 11 ++++++++ 7 files changed, 96 insertions(+), 7 deletions(-) diff --git a/django/core/management/base.py b/django/core/management/base.py index 7b31a69478..665b7da1fb 100644 --- a/django/core/management/base.py +++ b/django/core/management/base.py @@ -89,7 +89,7 @@ class OutputWrapper(object): @style_func.setter def style_func(self, style_func): - if style_func and hasattr(self._out, 'isatty') and self._out.isatty(): + if style_func and self.isatty(): self._style_func = style_func else: self._style_func = lambda x: x @@ -102,6 +102,9 @@ class OutputWrapper(object): def __getattr__(self, name): return getattr(self._out, name) + def isatty(self): + return hasattr(self._out, 'isatty') and self._out.isatty() + def write(self, msg, style_func=None, ending=None): ending = self.ending if ending is None else ending if ending and not msg.endswith(ending): diff --git a/django/core/management/commands/dumpdata.py b/django/core/management/commands/dumpdata.py index 01b229ce36..1468a92e6c 100644 --- a/django/core/management/commands/dumpdata.py +++ b/django/core/management/commands/dumpdata.py @@ -127,8 +127,11 @@ class Command(BaseCommand): raise CommandError("Unknown serialization format: %s" % format) - def get_objects(): - # Collate the objects to be serialized. + def get_objects(count_only=False): + """ + Collate the objects to be serialized. If count_only is True, just + count the number of objects to be serialized. + """ for model in serializers.sort_dependencies(app_list.items()): if model in excluded_models: continue @@ -141,17 +144,27 @@ class Command(BaseCommand): queryset = objects.using(using).order_by(model._meta.pk.name) if primary_keys: queryset = queryset.filter(pk__in=primary_keys) - for obj in queryset.iterator(): - yield obj + if count_only: + yield queryset.order_by().count() + else: + for obj in queryset.iterator(): + yield obj try: self.stdout.ending = None + progress_output = None + object_count = 0 + # If dumpdata is outputting to stdout, there is no way to display progress + if (output and self.stdout.isatty() and options['verbosity'] > 0): + progress_output = self.stdout + object_count = sum(get_objects(count_only=True)) stream = open(output, 'w') if output else None try: serializers.serialize(format, get_objects(), indent=indent, use_natural_foreign_keys=use_natural_foreign_keys, use_natural_primary_keys=use_natural_primary_keys, - stream=stream or self.stdout) + stream=stream or self.stdout, progress_output=progress_output, + object_count=object_count) finally: if stream: stream.close() diff --git a/django/core/serializers/base.py b/django/core/serializers/base.py index 23a0c7b72f..0e2e25d946 100644 --- a/django/core/serializers/base.py +++ b/django/core/serializers/base.py @@ -27,6 +27,29 @@ class DeserializationError(Exception): return cls("%s: (%s:pk=%s) field_value was '%s'" % (original_exc, model, fk, field_value)) +class ProgressBar(object): + progress_width = 75 + + def __init__(self, output, total_count): + self.output = output + self.total_count = total_count + self.prev_done = 0 + + def update(self, count): + if not self.output: + return + perc = count * 100 // self.total_count + done = perc * self.progress_width // 100 + if self.prev_done >= done: + return + self.prev_done = done + cr = '' if self.total_count == 1 else '\r' + self.output.write(cr + '[' + '.' * done + ' ' * (self.progress_width - done) + ']') + if done == self.progress_width: + self.output.write('\n') + self.output.flush() + + class Serializer(object): """ Abstract serializer base class. @@ -35,6 +58,7 @@ class Serializer(object): # Indicates if the implemented serializer is only available for # internal Django use. internal_use_only = False + progress_class = ProgressBar def serialize(self, queryset, **options): """ @@ -46,10 +70,13 @@ class Serializer(object): self.selected_fields = options.pop("fields", None) self.use_natural_foreign_keys = options.pop('use_natural_foreign_keys', False) self.use_natural_primary_keys = options.pop('use_natural_primary_keys', False) + progress_bar = self.progress_class( + options.pop('progress_output', None), options.pop('object_count', 0) + ) self.start_serialization() self.first = True - for obj in queryset: + for count, obj in enumerate(queryset, start=1): self.start_object(obj) # Use the concrete parent class' _meta instead of the object's _meta # This is to avoid local_fields problems for proxy models. Refs #17717. @@ -67,6 +94,7 @@ class Serializer(object): if self.selected_fields is None or field.attname in self.selected_fields: self.handle_m2m_field(obj, field) self.end_object(obj) + progress_bar.update(count) if self.first: self.first = False self.end_serialization() diff --git a/docs/ref/django-admin.txt b/docs/ref/django-admin.txt index 23c4de9fab..25b1a9a3f9 100644 --- a/docs/ref/django-admin.txt +++ b/docs/ref/django-admin.txt @@ -309,6 +309,12 @@ one model. By default ``dumpdata`` will output all the serialized data to standard output. This option allows you to specify the file to which the data is to be written. +When this option is set and the verbosity is greater than 0 (the default), a +progress bar is shown in the terminal. + +.. versionchanged:: 1.9 + + The progress bar in the terminal was added. flush ----- diff --git a/docs/releases/1.9.txt b/docs/releases/1.9.txt index 148e857c86..3c35ea23a2 100644 --- a/docs/releases/1.9.txt +++ b/docs/releases/1.9.txt @@ -364,6 +364,8 @@ Management Commands preceded by the operation's description. * The :djadmin:`dumpdata` command output is now deterministically ordered. + Moreover, when the ``--ouput`` option is specified, it also shows a progress + bar in the terminal. * The :djadmin:`createcachetable` command now has a ``--dry-run`` flag to print out the SQL rather than execute it. diff --git a/tests/fixtures/tests.py b/tests/fixtures/tests.py index ab86bb7b79..4c0c9ee88f 100644 --- a/tests/fixtures/tests.py +++ b/tests/fixtures/tests.py @@ -9,6 +9,7 @@ import warnings from django.apps import apps from django.contrib.sites.models import Site from django.core import management +from django.core.serializers.base import ProgressBar from django.db import IntegrityError, connection from django.test import ( TestCase, TransactionTestCase, ignore_warnings, skipUnlessDBFeature, @@ -286,6 +287,31 @@ class FixtureLoadingTests(DumpDataAssertMixin, TestCase): self._dumpdata_assert(['fixtures'], '[{"pk": 1, "model": "fixtures.category", "fields": {"description": "Latest news stories", "title": "News Stories"}}, {"pk": 2, "model": "fixtures.article", "fields": {"headline": "Poker has no place on ESPN", "pub_date": "2006-06-16T12:00:00"}}, {"pk": 3, "model": "fixtures.article", "fields": {"headline": "Time to reform copyright", "pub_date": "2006-06-16T13:00:00"}}]', filename='dumpdata.json') + def test_dumpdata_progressbar(self): + """ + Dumpdata shows a progress bar on the command line when --output is set, + stdout is a tty, and verbosity > 0. + """ + management.call_command('loaddata', 'fixture1.json', verbosity=0) + new_io = six.StringIO() + new_io.isatty = lambda: True + _, filename = tempfile.mkstemp() + options = { + 'format': 'json', + 'stdout': new_io, + 'stderr': new_io, + 'output': filename, + } + management.call_command('dumpdata', 'fixtures', **options) + self.assertTrue(new_io.getvalue().endswith('[' + '.' * ProgressBar.progress_width + ']\n')) + + # Test no progress bar when verbosity = 0 + options['verbosity'] = 0 + new_io = six.StringIO() + new_io.isatty = lambda: True + management.call_command('dumpdata', 'fixtures', **options) + self.assertEqual(new_io.getvalue(), '') + def test_compress_format_loading(self): # Load fixture 4 (compressed), using format specification management.call_command('loaddata', 'fixture4.json', verbosity=0) diff --git a/tests/serializers/tests.py b/tests/serializers/tests.py index 4703d59c0d..f6ebcba23a 100644 --- a/tests/serializers/tests.py +++ b/tests/serializers/tests.py @@ -9,6 +9,7 @@ from datetime import datetime from xml.dom import minidom from django.core import management, serializers +from django.core.serializers.base import ProgressBar from django.db import connection, transaction from django.test import ( SimpleTestCase, TestCase, TransactionTestCase, mock, override_settings, @@ -188,6 +189,16 @@ class SerializersTestBase(object): mv_obj = obj_list[0].object self.assertEqual(mv_obj.title, movie_title) + def test_serialize_progressbar(self): + fake_stdout = StringIO() + serializers.serialize( + self.serializer_name, Article.objects.all(), + progress_output=fake_stdout, object_count=Article.objects.count() + ) + self.assertTrue( + fake_stdout.getvalue().endswith('[' + '.' * ProgressBar.progress_width + ']\n') + ) + def test_serialize_superfluous_queries(self): """Ensure no superfluous queries are made when serializing ForeignKeys