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

@@ -89,7 +89,7 @@ def style_func(self):
 
     @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 @@ def __init__(self, out, style_func=None, ending='\n'):
     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):

django/core/management/commands/dumpdata.py

@@ -127,8 +127,11 @@ def handle(self, *app_labels, **options):
 
             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 @@ def get_objects():
                     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()

django/core/serializers/base.py

@@ -27,6 +27,29 @@ def WithData(cls, original_exc, model, fk, field_value):
         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 @@ def serialize(self, queryset, **options):
         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 @@ def serialize(self, queryset, **options):
                     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()

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

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.

tests/fixtures/tests.py

@@ -9,6 +9,7 @@
 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 @@ def test_dumpdata_with_file_output(self):
         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)

tests/serializers/tests.py

@@ -9,6 +9,7 @@
 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 @@ def test_serialize_unicode(self):
         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