Add progress bars to several commands and plugins.

Many long-running commands produce little or no feedback in the terminal to
indicate that they're progressing, and none of them provide estimates of how
long the operation will run. This change introduces the `enlighten` python
package, which displays progress bars akin to TQDM below the existing terminal
output.

To support consistent use and presentation of the progress bars, and to
allow for future modification, we introduce a method to beets.ui -
beets.ui.iprogress_bar - which can be used by Beets' core commands and all
Beets plugins. Example usage is provided in the methods' documentation
and in a new 'further reading' doc for developers.

The Enlighten library does not work as well in Windows PowerShell as it
does in a linux terminal (manually tested in Zsh), so the progress bars
are disabled in Windows environments. Resolving these issues and enabling
them in Windows is left as future work.

Progress Bars are disabled also when not in a TTY, to prevent
interference with logging when running as a web server, for example.

Author: peterjdolan

beets/ui/__init__.py

@@ -28,22 +28,41 @@
 import sys
 import textwrap
 import traceback
+import typing
 import warnings
 from difflib import SequenceMatcher
 from functools import cache
 from itertools import chain
-from typing import Any, Callable, Literal
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Generator,
+    Iterable,
+    Literal,
+    Sized,
+    TypeVar,
+    Union,
+)
 
 import confuse
+import enlighten
+from typing_extensions import Protocol
 
 from beets import config, library, logging, plugins, util
 from beets.dbcore import db
 from beets.dbcore import query as db_query
 from beets.util import as_string
 from beets.util.functemplate import template
 
+if TYPE_CHECKING:
+    from beets.db import Album as Album
+    from beets.db import Item as Item
+
+is_windows = sys.platform == "win32"
+
 # On Windows platforms, use colorama to support "ANSI" terminal colors.
-if sys.platform == "win32":
+if is_windows:
     try:
         import colorama
     except ImportError:
@@ -1325,6 +1344,80 @@ def add_all_common_options(self):
         self.add_format_option()
 
 
+T = TypeVar("T", covariant=True)
+U = TypeVar("U")
+
+
+class SizedIterable(Protocol[T], Sized, Iterable[T]):
+    pass
+
+
+def iprogress_bar(
+    sequence: Union[Iterable[U], SizedIterable[U]], **kwargs
+) -> Generator[U, None, None]:
+    """Construct and manage an `enlighten.Counter` progress bar while iterating.
+
+    Example usage:
+    ```
+    for album in ui.iprogress_bar(
+        lib.albums(), desc="Updating albums", unit="albums"):
+        do_something_to(album)
+    ```
+
+    If the progress bar is iterating over an Album or an Item, then it will detect
+    whether or not the item has been modified, and will color-code the progress bar
+    with white and blue to indicate total progress and the portion of items that have
+    been modified.
+
+    Args:
+        sequence: An `Iterable` sequence to iterate over. If provided, and the
+            sequence can return its length, then the length will be used as the
+            total for the counter. The counter will be updated for each item
+            in the sequence.
+        kwargs: Additional keyword arguments to pass to the `enlighten.Counter`
+            constructor.
+
+    Yields:
+        The items from the sequence.
+    """
+    if sequence is None:
+        log.error("sequence must not be None")
+        return
+
+    # If the total was not directly set, and the iterable is sized, then use its size as
+    # the progress bar's total.
+    if "total" not in kwargs:
+        if hasattr(sequence, "__len__"):
+            sized_seq = typing.cast(SizedIterable[U], sequence)
+            kwargs["total"] = len(sized_seq)
+
+    # Disabled in windows environments and when not attached to a TTY. See method docs
+    # for details.
+    with enlighten.Manager(
+        enabled=not is_windows
+        and (hasattr(sys.stdout, "isatty") and sys.stdout.isatty())
+    ) as manager:
+        with manager.counter(**kwargs) as counter:
+            change_counter = counter.add_subcounter("blue")
+
+            for item in sequence:
+                revision = None
+                if hasattr(item, "_revision"):
+                    revision = item._revision
+
+                # Yield the item, allowing it to be modified, or not.
+                yield item
+
+                if (
+                    revision
+                    and hasattr(item, "_revision")
+                    and item._revision != revision
+                ):
+                    change_counter.update()
+                else:
+                    counter.update()
+
+
 # Subcommand parsing infrastructure.
 #
 # This is a fairly generic subcommand parser for optparse. It is

beetsplug/autobpm.py

@@ -21,7 +21,7 @@
 import numpy as np
 
 from beets.plugins import BeetsPlugin
-from beets.ui import Subcommand, should_write
+from beets.ui import Subcommand, iprogress_bar, should_write
 
 if TYPE_CHECKING:
     from beets.importer import ImportTask
@@ -56,7 +56,7 @@ def imported(self, _, task: ImportTask) -> None:
         self.calculate_bpm(task.imported_items())
 
     def calculate_bpm(self, items: list[Item], write: bool = False) -> None:
-        for item in items:
+        for item in iprogress_bar(items, desc="Calculating BPM", unit="items"):
             path = item.filepath
             if bpm := item.bpm:
                 self._log.info("BPM for {} already exists: {}", path, bpm)

beetsplug/badfiles.py

@@ -14,18 +14,20 @@
 
 """Use command-line tools to check for audio file corruption."""
 
+import concurrent.futures
 import errno
 import os
 import shlex
 import sys
 from subprocess import STDOUT, CalledProcessError, check_output, list2cmdline
+from typing import Callable, Optional, Union
 
 import confuse
 
-from beets import importer, ui
+from beets import importer, library, ui
 from beets.plugins import BeetsPlugin
 from beets.ui import Subcommand
-from beets.util import displayable_path, par_map
+from beets.util import displayable_path
 
 
 class CheckerCommandError(Exception):
@@ -45,6 +47,13 @@ def __init__(self, cmd, oserror):
         self.msg = str(oserror)
 
 
+# CheckResult is a tuple of 1. status code, 2. how many errors there were, and 3.
+# a list of error output messages.
+CheckResult = tuple[int, int, list[str]]
+
+CheckMethod = Callable[[str], CheckResult]
+
+
 class BadFiles(BeetsPlugin):
     def __init__(self):
         super().__init__()
@@ -55,12 +64,12 @@ def __init__(self):
             "import_task_before_choice", self.on_import_task_before_choice
         )
 
-    def run_command(self, cmd):
+    def run_command(self, cmd: list[str]) -> CheckResult:
         self._log.debug(
             "running command: {}", displayable_path(list2cmdline(cmd))
         )
         try:
-            output = check_output(cmd, stderr=STDOUT)
+            output: bytes = check_output(cmd, stderr=STDOUT)
             errors = 0
             status = 0
         except CalledProcessError as e:
@@ -69,28 +78,28 @@ def run_command(self, cmd):
             status = e.returncode
         except OSError as e:
             raise CheckerCommandError(cmd, e)
-        output = output.decode(sys.getdefaultencoding(), "replace")
-        return status, errors, [line for line in output.split("\n") if line]
+        output_dec = output.decode(sys.getdefaultencoding(), "replace")
+        return status, errors, [line for line in output_dec.split("\n") if line]
 
-    def check_mp3val(self, path):
+    def check_mp3val(self, path: str) -> CheckResult:
         status, errors, output = self.run_command(["mp3val", path])
         if status == 0:
             output = [line for line in output if line.startswith("WARNING:")]
             errors = len(output)
         return status, errors, output
 
-    def check_flac(self, path):
+    def check_flac(self, path: str) -> CheckResult:
         return self.run_command(["flac", "-wst", path])
 
-    def check_custom(self, command):
+    def check_custom(self, command: str) -> Callable[[str], CheckResult]:
         def checker(path):
             cmd = shlex.split(command)
             cmd.append(path)
             return self.run_command(cmd)
 
         return checker
 
-    def get_checker(self, ext):
+    def get_checker(self, ext: str) -> Optional[CheckMethod]:
         ext = ext.lower()
         try:
             command = self.config["commands"].get(dict).get(ext)
@@ -102,8 +111,9 @@ def get_checker(self, ext):
             return self.check_mp3val
         if ext == "flac":
             return self.check_flac
+        return None
 
-    def check_item(self, item):
+    def check_item(self, item: library.Item) -> tuple[bool, list[str]]:
         # First, check whether the path exists. If not, the user
         # should probably run `beet update` to cleanup your library.
         dpath = displayable_path(item.path)
@@ -118,8 +128,8 @@ def check_item(self, item):
         checker = self.get_checker(ext)
         if not checker:
             self._log.error("no checker specified in the config for {}", ext)
-            return []
-        path = item.path
+            return False, []
+        path: Union[bytes, str] = item.path
         if not isinstance(path, str):
             path = item.path.decode(sys.getfilesystemencoding())
         try:
@@ -132,11 +142,13 @@ def check_item(self, item):
                 )
             else:
                 self._log.error("error invoking {0.checker}: {0.msg}", e)
-            return []
+            return False, []
 
+        success = True
         error_lines = []
 
         if status > 0:
+            success = False
             error_lines.append(
                 f"{ui.colorize('text_error', dpath)}: checker exited with"
                 f" status {status}"
@@ -145,6 +157,7 @@ def check_item(self, item):
                 error_lines.append(f"  {line}")
 
         elif errors > 0:
+            success = False
             error_lines.append(
                 f"{ui.colorize('text_warning', dpath)}: checker found"
                 f" {status} errors or warnings"
@@ -154,23 +167,24 @@ def check_item(self, item):
         elif self.verbose:
             error_lines.append(f"{ui.colorize('text_success', dpath)}: ok")
 
-        return error_lines
+        return success, error_lines
 
-    def on_import_task_start(self, task, session):
+    def on_import_task_start(self, task, session) -> None:
         if not self.config["check_on_import"].get(False):
             return
 
         checks_failed = []
 
         for item in task.items:
-            error_lines = self.check_item(item)
-            if error_lines:
-                checks_failed.append(error_lines)
+            _, error_lines = self.check_item(item)
+            checks_failed.append(error_lines)
 
         if checks_failed:
             task._badfiles_checks_failed = checks_failed
 
-    def on_import_task_before_choice(self, task, session):
+    def on_import_task_before_choice(
+        self, task, session
+    ) -> Optional[importer.Action]:
         if hasattr(task, "_badfiles_checks_failed"):
             ui.print_(
                 f"{ui.colorize('text_warning', 'BAD')} one or more files failed"
@@ -194,16 +208,27 @@ def on_import_task_before_choice(self, task, session):
             else:
                 raise Exception(f"Unexpected selection: {sel}")
 
-    def command(self, lib, opts, args):
+        return None
+
+    def command(self, lib, opts, args) -> None:
         # Get items from arguments
         items = lib.items(args)
         self.verbose = opts.verbose
 
         def check_and_print(item):
-            for error_line in self.check_item(item):
-                ui.print_(error_line)
-
-        par_map(check_and_print, items)
+            success, error_lines = self.check_item(item)
+            if not success:
+                for line in error_lines:
+                    ui.print_(line)
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            for _ in ui.iprogress_bar(
+                executor.map(check_and_print, items),
+                desc="Checking",
+                unit="item",
+                total=len(items),
+            ):
+                pass
 
     def commands(self):
         bad_command = Subcommand(

beetsplug/chroma.py

@@ -254,7 +254,11 @@ def submit_cmd_func(lib, opts, args):
         )
 
         def fingerprint_cmd_func(lib, opts, args):
-            for item in lib.items(args):
+            for item in ui.iprogress_bar(
+                lib.items(args),
+                desc="Fingerprinting items",
+                unit="items",
+            ):
                 fingerprint_item(self._log, item, write=ui.should_write())
 
         fingerprint_cmd.func = fingerprint_cmd_func

beetsplug/duplicates.py

@@ -17,6 +17,7 @@
 import os
 import shlex
 
+from beets import ui
 from beets.library import Album, Item
 from beets.plugins import BeetsPlugin
 from beets.ui import Subcommand, UserError, print_
@@ -284,9 +285,13 @@ def _group_by(self, objs, keys, strict):
         import collections
 
         counts = collections.defaultdict(list)
-        for obj in objs:
+        for obj in ui.iprogress_bar(
+            objs,
+            desc="Finding duplicates",
+            unit="items",
+        ):
             values = [getattr(obj, k, None) for k in keys]
-            values = [v for v in values if v not in (None, "")]
+            values = list(filter(lambda v: v not in (None, ""), values))
             if strict and len(values) < len(keys):
                 self._log.debug(
                     "some keys {} on item {.filepath} are null or empty: skipping",