Merge pull request #1449 from enkore/f/extractprogress

extract --progress

9 files changed, 74 insertions, 42 deletions

diff --git a/src/borg/archive.py b/src/borg/archive.py
index 8548e3aa3..b17685429 100644
--- a/src/borg/archive.py
+++ b/src/borg/archive.py
@@ -422,7 +422,7 @@ Number of files: {0.stats.nfiles}'''.format(
         return stats
 
     def extract_item(self, item, restore_attrs=True, dry_run=False, stdout=False, sparse=False,
-                     hardlink_masters=None, original_path=None):
+                     hardlink_masters=None, original_path=None, pi=None):
         """
         Extract archive item.
 
@@ -433,11 +433,14 @@ Number of files: {0.stats.nfiles}'''.format(
         :param sparse: write sparse files (chunk-granularity, independent of the original being sparse)
         :param hardlink_masters: maps paths to (chunks, link_target) for extracting subtrees with hardlinks correctly
         :param original_path: 'path' key as stored in archive
+        :param pi: ProgressIndicatorPercent (or similar) for file extraction progress (in bytes)
         """
         has_damaged_chunks = 'chunks_healthy' in item
         if dry_run or stdout:
             if 'chunks' in item:
                 for _, data in self.pipeline.fetch_many([c.id for c in item.chunks], is_preloaded=True):
+                    if pi:
+                        pi.show(increase=len(data))
                     if stdout:
                         sys.stdout.buffer.write(data)
                 if stdout:
@@ -489,6 +492,8 @@ Number of files: {0.stats.nfiles}'''.format(
             with fd:
                 ids = [c.id for c in item.chunks]
                 for _, data in self.pipeline.fetch_many(ids, is_preloaded=True):
+                    if pi:
+                        pi.show(increase=len(data))
                     with backup_io():
                         if sparse and self.zeros.startswith(data):
                             # all-zero chunk: create a hole in a sparse file
@@ -635,7 +640,7 @@ Number of files: {0.stats.nfiles}'''.format(
         try:
             unpacker = msgpack.Unpacker(use_list=False)
             items_ids = self.metadata.items
-            pi = ProgressIndicatorPercent(total=len(items_ids), msg="Decrementing references %3.0f%%", same_line=True)
+            pi = ProgressIndicatorPercent(total=len(items_ids), msg="Decrementing references %3.0f%%")
             for (i, (items_id, data)) in enumerate(zip(items_ids, self.repository.get_many(items_ids))):
                 if progress:
                     pi.show(i)
@@ -1028,7 +1033,7 @@ class ArchiveChecker:
         logger.info('Starting cryptographic data integrity verification...')
         count = len(self.chunks)
         errors = 0
-        pi = ProgressIndicatorPercent(total=count, msg="Verifying data %6.2f%%", step=0.01, same_line=True)
+        pi = ProgressIndicatorPercent(total=count, msg="Verifying data %6.2f%%", step=0.01)
         for chunk_id, (refcount, *_) in self.chunks.iteritems():
             pi.show()
             try:
diff --git a/src/borg/archiver.py b/src/borg/archiver.py
index 9074f8994..619dbd7e9 100644
--- a/src/borg/archiver.py
+++ b/src/borg/archiver.py
@@ -41,6 +41,7 @@ from .helpers import log_multi
 from .helpers import parse_pattern, PatternMatcher, PathPrefixPattern
 from .helpers import signal_handler
 from .helpers import ErrorIgnoringTextIOWrapper
+from .helpers import ProgressIndicatorPercent
 from .item import Item
 from .key import key_creator, RepoKey, PassphraseKey
 from .platform import get_flags
@@ -439,6 +440,7 @@ class Archiver:
 
         matcher, include_patterns = self.build_matcher(args.excludes, args.paths)
 
+        progress = args.progress
         output_list = args.output_list
         dry_run = args.dry_run
         stdout = args.stdout
@@ -453,6 +455,14 @@ class Archiver:
                     item.get('hardlink_master', True) and 'source' not in item)
 
         filter = self.build_filter(matcher, item_is_hardlink_master, strip_components)
+        if progress:
+            progress_logger = logging.getLogger(ProgressIndicatorPercent.LOGGER)
+            progress_logger.info('Calculating size')
+            extracted_size = sum(item.file_size() for item in archive.iter_items(filter))
+            pi = ProgressIndicatorPercent(total=extracted_size, msg='Extracting files %5.1f%%', step=0.1)
+        else:
+            pi = None
+
         for item in archive.iter_items(filter, preload=True):
             orig_path = item.path
             if item_is_hardlink_master(item):
@@ -472,19 +482,21 @@ class Archiver:
                 logging.getLogger('borg.output.list').info(remove_surrogates(orig_path))
             try:
                 if dry_run:
-                    archive.extract_item(item, dry_run=True)
+                    archive.extract_item(item, dry_run=True, pi=pi)
                 else:
                     if stat.S_ISDIR(item.mode):
                         dirs.append(item)
                         archive.extract_item(item, restore_attrs=False)
                     else:
                         archive.extract_item(item, stdout=stdout, sparse=sparse, hardlink_masters=hardlink_masters,
-                                             original_path=orig_path)
+                                             original_path=orig_path, pi=pi)
             except BackupOSError as e:
                 self.print_warning('%s: %s', remove_surrogates(orig_path), e)
 
         if not args.dry_run:
+            pi = ProgressIndicatorPercent(total=len(dirs), msg='Setting directory permissions %3.0f%%')
             while dirs:
+                pi.show()
                 dir_item = dirs.pop(-1)
                 try:
                     archive.extract_item(dir_item)
@@ -1641,6 +1653,9 @@ class Archiver:
                                           formatter_class=argparse.RawDescriptionHelpFormatter,
                                           help='extract archive contents')
         subparser.set_defaults(func=self.do_extract)
+        subparser.add_argument('-p', '--progress', dest='progress',
+                               action='store_true', default=False,
+                               help='show progress while extracting (may be slower)')
         subparser.add_argument('--list', dest='output_list',
                                action='store_true', default=False,
                                help='output verbose list of items (files, dirs, ...)')
diff --git a/src/borg/helpers.py b/src/borg/helpers.py
index 6751285d8..3d30692f2 100644
--- a/src/borg/helpers.py
+++ b/src/borg/helpers.py
@@ -1105,14 +1105,15 @@ def yes(msg=None, false_msg=None, true_msg=None, default_msg=None,
 
 
 class ProgressIndicatorPercent:
-    def __init__(self, total, step=5, start=0, same_line=False, msg="%3.0f%%"):
+    LOGGER = 'borg.output.progress'
+
+    def __init__(self, total, step=5, start=0, msg="%3.0f%%"):
         """
         Percentage-based progress indicator
 
         :param total: total amount of items
         :param step: step size in percent
         :param start: at which percent value to start
-        :param same_line: if True, emit output always on same line
         :param msg: output message, must contain one %f placeholder for the percentage
         """
         self.counter = 0  # 0 .. (total-1)
@@ -1120,9 +1121,8 @@ class ProgressIndicatorPercent:
         self.trigger_at = start  # output next percentage value when reaching (at least) this
         self.step = step
         self.msg = msg
-        self.same_line = same_line
         self.handler = None
-        self.logger = logging.getLogger('borg.output.progress')
+        self.logger = logging.getLogger(self.LOGGER)
 
         # If there are no handlers, set one up explicitly because the
         # terminator and propagation needs to be set.  If there are,
@@ -1130,7 +1130,7 @@ class ProgressIndicatorPercent:
         if not self.logger.handlers:
             self.handler = logging.StreamHandler(stream=sys.stderr)
             self.handler.setLevel(logging.INFO)
-            self.handler.terminator = '\r' if self.same_line else '\n'
+            self.handler.terminator = '\r'
 
             self.logger.addHandler(self.handler)
             if self.logger.level == logging.NOTSET:
@@ -1142,17 +1142,17 @@ class ProgressIndicatorPercent:
             self.logger.removeHandler(self.handler)
             self.handler.close()
 
-    def progress(self, current=None):
+    def progress(self, current=None, increase=1):
         if current is not None:
             self.counter = current
         pct = self.counter * 100 / self.total
-        self.counter += 1
+        self.counter += increase
         if pct >= self.trigger_at:
             self.trigger_at += self.step
             return pct
 
-    def show(self, current=None):
-        pct = self.progress(current)
+    def show(self, current=None, increase=1):
+        pct = self.progress(current, increase)
         if pct is not None:
             return self.output(pct)
 
@@ -1160,8 +1160,7 @@ class ProgressIndicatorPercent:
         self.logger.info(self.msg % percent)
 
     def finish(self):
-        if self.same_line:
-            self.logger.info(" " * len(self.msg % 100.0))
+        self.logger.info(" " * len(self.msg % 100.0))
 
 
 class ProgressIndicatorEndless:
diff --git a/src/borg/item.py b/src/borg/item.py
index d74bfdb55..0bc336239 100644
--- a/src/borg/item.py
+++ b/src/borg/item.py
@@ -157,6 +157,11 @@ class Item(PropDict):
 
     part = PropDict._make_property('part', int)
 
+    def file_size(self):
+        if 'chunks' not in self:
+            return 0
+        return sum(chunk.size for chunk in self.chunks)
+
 
 class EncryptedKey(PropDict):
     """
diff --git a/src/borg/repository.py b/src/borg/repository.py
index 9ae7e36f8..8b5f85db5 100644
--- a/src/borg/repository.py
+++ b/src/borg/repository.py
@@ -536,7 +536,7 @@ class Repository:
         self.prepare_txn(index_transaction_id, do_cleanup=False)
         try:
             segment_count = sum(1 for _ in self.io.segment_iterator())
-            pi = ProgressIndicatorPercent(total=segment_count, msg="Replaying segments %3.0f%%", same_line=True)
+            pi = ProgressIndicatorPercent(total=segment_count, msg="Replaying segments %3.0f%%")
             for i, (segment, filename) in enumerate(self.io.segment_iterator()):
                 pi.show(i)
                 if index_transaction_id is not None and segment <= index_transaction_id:
@@ -636,7 +636,7 @@ class Repository:
         segments_transaction_id = self.io.get_segments_transaction_id()
         self.prepare_txn(None)  # self.index, self.compact, self.segments all empty now!
         segment_count = sum(1 for _ in self.io.segment_iterator())
-        pi = ProgressIndicatorPercent(total=segment_count, msg="Checking segments %3.1f%%", step=0.1, same_line=True)
+        pi = ProgressIndicatorPercent(total=segment_count, msg="Checking segments %3.1f%%", step=0.1)
         for i, (segment, filename) in enumerate(self.io.segment_iterator()):
             pi.show(i)
             if segment > transaction_id:
diff --git a/src/borg/testsuite/archiver.py b/src/borg/testsuite/archiver.py
index 22438da0e..45daceda0 100644
--- a/src/borg/testsuite/archiver.py
+++ b/src/borg/testsuite/archiver.py
@@ -751,6 +751,15 @@ class ArchiverTestCase(ArchiverTestCaseBase):
             output = self.cmd('extract', '--list', '--info', self.repository_location + '::test')
         self.assert_in("input/file", output)
 
+    def test_extract_progress(self):
+        self.cmd('init', self.repository_location)
+        self.create_regular_file('file', size=1024 * 80)
+        self.cmd('create', self.repository_location + '::test', 'input')
+
+        with changedir('output'):
+            output = self.cmd('extract', self.repository_location + '::test', '--progress')
+            assert 'Extracting files' in output
+
     def _create_test_caches(self):
         self.cmd('init', self.repository_location)
         self.create_regular_file('file1', size=1024 * 80)
diff --git a/src/borg/testsuite/helpers.py b/src/borg/testsuite/helpers.py
index 0f8a853e0..5f5a3806c 100644
--- a/src/borg/testsuite/helpers.py
+++ b/src/borg/testsuite/helpers.py
@@ -870,24 +870,9 @@ def test_yes_env_output(capfd, monkeypatch):
     assert 'yes' in err
 
 
-def test_progress_percentage_multiline(capfd):
-    pi = ProgressIndicatorPercent(1000, step=5, start=0, same_line=False, msg="%3.0f%%")
-    pi.show(0)
-    out, err = capfd.readouterr()
-    assert err == '  0%\n'
-    pi.show(420)
-    out, err = capfd.readouterr()
-    assert err == ' 42%\n'
-    pi.show(1000)
-    out, err = capfd.readouterr()
-    assert err == '100%\n'
-    pi.finish()
-    out, err = capfd.readouterr()
-    assert err == ''
-
-
 def test_progress_percentage_sameline(capfd):
-    pi = ProgressIndicatorPercent(1000, step=5, start=0, same_line=True, msg="%3.0f%%")
+    pi = ProgressIndicatorPercent(1000, step=5, start=0, msg="%3.0f%%")
+    pi.logger.setLevel('INFO')
     pi.show(0)
     out, err = capfd.readouterr()
     assert err == '  0%\r'
@@ -904,22 +889,22 @@ def test_progress_percentage_sameline(capfd):
 
 
 def test_progress_percentage_step(capfd):
-    pi = ProgressIndicatorPercent(100, step=2, start=0, same_line=False, msg="%3.0f%%")
+    pi = ProgressIndicatorPercent(100, step=2, start=0, msg="%3.0f%%")
+    pi.logger.setLevel('INFO')
     pi.show()
     out, err = capfd.readouterr()
-    assert err == '  0%\n'
+    assert err == '  0%\r'
     pi.show()
     out, err = capfd.readouterr()
     assert err == ''  # no output at 1% as we have step == 2
     pi.show()
     out, err = capfd.readouterr()
-    assert err == '  2%\n'
+    assert err == '  2%\r'
 
 
 def test_progress_percentage_quiet(capfd):
-    logging.getLogger('borg.output.progress').setLevel(logging.WARN)
-
-    pi = ProgressIndicatorPercent(1000, step=5, start=0, same_line=False, msg="%3.0f%%")
+    pi = ProgressIndicatorPercent(1000, step=5, start=0, msg="%3.0f%%")
+    pi.logger.setLevel('WARN')
     pi.show(0)
     out, err = capfd.readouterr()
     assert err == ''
diff --git a/src/borg/testsuite/item.py b/src/borg/testsuite/item.py
index b0b7569e3..fc60e91df 100644
--- a/src/borg/testsuite/item.py
+++ b/src/borg/testsuite/item.py
@@ -1,5 +1,6 @@
 import pytest
 
+from ..cache import ChunkListEntry
 from ..item import Item
 from ..helpers import StableDict
 
@@ -145,3 +146,16 @@ def test_unknown_property():
     item = Item()
     with pytest.raises(AttributeError):
         item.unknown_attribute = None
+
+
+def test_item_file_size():
+    item = Item(chunks=[
+        ChunkListEntry(csize=1, size=1000, id=None),
+        ChunkListEntry(csize=1, size=2000, id=None),
+    ])
+    assert item.file_size() == 3000
+
+
+def test_item_file_size_no_chunks():
+    item = Item()
+    assert item.file_size() == 0
diff --git a/src/borg/upgrader.py b/src/borg/upgrader.py
index 42f9a4696..69712832a 100644
--- a/src/borg/upgrader.py
+++ b/src/borg/upgrader.py
@@ -77,7 +77,7 @@ class AtticRepositoryUpgrader(Repository):
         replace the 8 first bytes of all regular files in there."""
         logger.info("converting %d segments..." % len(segments))
         segment_count = len(segments)
-        pi = ProgressIndicatorPercent(total=segment_count, msg="Converting segments %3.0f%%", same_line=True)
+        pi = ProgressIndicatorPercent(total=segment_count, msg="Converting segments %3.0f%%")
         for i, filename in enumerate(segments):
             if progress:
                 pi.show(i)