Fixed #22258 -- Added progress status for dumpdata when outputting to file

Thanks Gwildor Sok for the report and Tim Graham for the review.
This commit is contained in:
Claude Paroz 2015-07-21 23:24:32 +02:00
parent 03aec35a12
commit c296e55dc6
7 changed files with 96 additions and 7 deletions

View File

@ -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):

View File

@ -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()

View File

@ -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()

View File

@ -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
-----

View File

@ -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.

View File

@ -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)

View File

@ -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