Add a progress bar for rsync

On a terminal, run rsync with --info=progress2 and draw its updates as a
progress bar on stderr. The ETA is estimated locally, since rsync reports
the elapsed time when a file finishes. Disable it with --no-progress.

simple_backup/simple_backup.py

@@ -23,7 +23,7 @@ import re
 from typing import Callable, Iterator, List, Optional, ParamSpec, TypedDict, TypeVar
 import warnings
 from functools import wraps
-from shutil import rmtree, which
+from shutil import get_terminal_size, rmtree, which
 import shlex
 import argparse
 import configparser
@@ -74,6 +74,10 @@ R = TypeVar('R')
 # Seconds to wait for the TCP connection to the SSH server
 SSH_TIMEOUT = 15
 
+# An rsync --info=progress2 update: size, percent, speed, time, then optionally (xfr#N, ir-chk/to-chk=N/N)
+PROGRESS_RE = re.compile(r'\s*(\S+)\s+(\d+)%\s+(\S+)\s+\d+:\d\d:\d\d(?:\s+\(xfr#\d+, (ir|to)-chk=\d+/\d+\))?\s*')
+PROGRESS_BAR_WIDTH = 30
+
 
 def timing(func: Callable[P, R]) -> Callable[P, R]:
     """Decorator to measure execution time of a function
@@ -100,6 +104,59 @@ class MyFormatter(argparse.RawTextHelpFormatter, argparse.ArgumentDefaultsHelpFo
     """Custom format for argparse help text"""
 
 
+class _ProgressBar:
+    """Draws rsync --info=progress2 updates as a bar on stderr"""
+
+    def __init__(self) -> None:
+        self._start = default_timer()
+        # rsync reports ir-chk while incremental recursion is still finding files
+        self._scanning = True
+        self._shown = False
+
+    def update(self, record: str) -> bool:
+        """Draw the bar if record is a progress update, and return whether it was"""
+
+        m = PROGRESS_RE.fullmatch(record)
+
+        if m is None:
+            return False
+
+        size, percent, speed, chk = m.group(1), int(m.group(2)), m.group(3), m.group(4)
+
+        if chk is not None:
+            self._scanning = chk == 'ir'
+
+        # rsync's own time field is the elapsed time on records ending a file, so estimate the ETA here
+        if 0 < percent < 100:
+            eta = int((default_timer() - self._start) * (100 - percent) / percent)
+            text = f'{percent:3d}%  {size}  {speed}  ETA {eta // 3600}:{eta // 60 % 60:02d}:{eta % 60:02d}'
+        else:
+            text = f'{percent:3d}%  {size}  {speed}'
+
+        if self._scanning:
+            text += '  (scanning)'
+
+        width = get_terminal_size().columns - 1
+
+        # Fixed width, so the bar doesn't jump when the text changes length
+        if len(text) + PROGRESS_BAR_WIDTH + 4 <= width:
+            filled = PROGRESS_BAR_WIDTH * percent // 100
+            text = f'[{"#" * filled}{"-" * (PROGRESS_BAR_WIDTH - filled)}]  {text}'
+
+        sys.stderr.write(f'\r{text[:width]}\033[K')
+        sys.stderr.flush()
+        self._shown = True
+
+        return True
+
+    def clear(self) -> None:
+        """Remove the bar, so that log messages start on a clean line"""
+
+        if self._shown:
+            sys.stderr.write('\r\033[K')
+            sys.stderr.flush()
+
+
 class Backup:
     """Main class defining parameters and functions for performing backup
 
@@ -124,6 +181,8 @@ class Backup:
             Run remote rsync with sudo
         remove_before: bool
             Indicate if removing old backups will be performed before copying files
+        progress: bool
+            Show a progress bar while rsync runs
 
     Methods:
         check_params():
@@ -141,7 +200,8 @@ class Backup:
     def __init__(self, inputs: Optional[List[str]], output: Optional[str], exclude: Optional[List[str]],
                  keep: Optional[int], options: str,
                  ssh_host: Optional[str] = None, ssh_user: Optional[str] = None, ssh_keyfile: Optional[str] = None,
-                 remote_sudo: bool = False, remove_before: bool = False, verbose: bool = False) -> None:
+                 remote_sudo: bool = False, remove_before: bool = False, verbose: bool = False,
+                 progress: bool = False) -> None:
         # Missing values are reported by check_params()
         self.inputs = inputs if inputs is not None else []
         self.output = output if output is not None else ''
@@ -154,6 +214,7 @@ class Backup:
         self.remote_sudo = remote_sudo
         self._remove_before = remove_before
         self._verbose = verbose
+        self._progress = progress
         self._last_backup = ''
         self._server = ''
         self._output_dir = ''
@@ -648,11 +709,25 @@ class Backup:
         if self._remote and self.remote_sudo:
             args.append('--rsync-path=sudo rsync')
 
+        bar = None
+
+        if self._progress:
+            args.append('--info=progress2')
+            bar = _ProgressBar()
+
         with Popen(args, stdin=PIPE, stdout=PIPE, stderr=STDOUT, shell=False) as p:
             assert p.stdin is not None and isinstance(p.stdout, io.BufferedIOBase)
             p.stdin.close()
 
-            output = [r for r in _read_records(p.stdout) if r != '']
+            output = []
+
+            try:
+                for r in _read_records(p.stdout):
+                    if r != '' and (bar is None or not bar.update(r)):
+                        output.append(r)
+            finally:
+                if bar is not None:
+                    bar.clear()
 
             try:
                 del os.environ['SSHPASS']
@@ -766,6 +841,8 @@ def _parse_arguments() -> argparse.Namespace:
                         help='Remove old backups before executing the backup, instead of after')
     parser.add_argument('--no-syslog', action='store_true',
                         help='Disable systemd journal logging')
+    parser.add_argument('--no-progress', action='store_true',
+                        help='Don\'t show a progress bar (it is only shown on a terminal)')
     parser.add_argument('--rsync-options', nargs='+',
                         choices=['a', 'l', 'p', 't', 'g', 'o',
                                  'c', 'h', 's', 'D', 'H', 'X'],
@@ -1039,7 +1116,8 @@ def simple_backup() -> int:
         rsync_options.append('--numeric-ids')
 
     backup = Backup(inputs, output, exclude, keep, ' '.join(rsync_options), ssh_host, ssh_user, ssh_keyfile,
-                    remote_sudo, remove_before=args.remove_before_backup, verbose=args.verbose)
+                    remote_sudo, remove_before=args.remove_before_backup, verbose=args.verbose,
+                    progress=not args.no_progress and sys.stderr.isatty())
 
     try:
         return_code = backup.check_params(homedir)