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 two methods to beets.ui -
beets.ui.iprogress_bar, and beets.ui.changed_unchanged_error_pbars -
which can be used by Beets' core commands and all Beets plugins. Example
usage is provided in the methods' documentation.

Integrating progress bars into the 'import' command required a more
custom implementation than the beets.ui.* commands could support. The
approach taken is certainly open to discussion.

Notably, 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.

Author: peterjdolan

beets/importer/tasks.py

@@ -234,7 +234,7 @@ def chosen_info(self):
             return likelies
         elif self.choice_flag is Action.APPLY and self.match:
             return self.match.info.copy()
-        assert False
+        raise ValueError("Invalid choice flag; this should never happen.")
 
     def imported_items(self):
         """Return a list of Items that should be added to the library.
@@ -249,7 +249,7 @@ def imported_items(self):
         ):
             return list(self.match.mapping.keys())
         else:
-            assert False
+            raise ValueError("Invalid choice flag; this should never happen.")
 
     def apply_metadata(self):
         """Copy metadata from match info to the items."""
@@ -318,6 +318,8 @@ def finalize(self, session: ImportSession):
         if not self.skip:
             self._emit_imported(session.lib)
 
+        session.task_finalized()
+
     def cleanup(self, copy=False, delete=False, move=False):
         """Remove and prune imported paths."""
         # Do not delete any files or prune directories when skipping.
@@ -355,9 +357,10 @@ def handle_created(self, session: ImportSession):
         else:
             # The plugins gave us a list of lists of tasks. Flatten it.
             tasks = [t for inner in tasks for t in inner]
+        session.tasks_created(tasks)
         return tasks
 
-    def lookup_candidates(self, search_ids: list[str]) -> None:
+    def lookup_candidates(self, session: ImportSession, search_ids: list[str]) -> None:
         """Retrieve and store candidates for this album.
 
         If User-specified ``search_ids`` list is not empty, the lookup is
@@ -366,6 +369,7 @@ def lookup_candidates(self, search_ids: list[str]) -> None:
         self.cur_artist, self.cur_album, (self.candidates, self.rec) = (
             autotag.tag_album(self.items, search_ids=search_ids)
         )
+        session.task_candidates_found()
 
     def find_duplicates(self, lib: library.Library) -> list[library.Album]:
         """Return a list of albums from `lib` with the same artist and
@@ -638,6 +642,7 @@ def choose_match(self, session):
         choice = session.choose_match(self)
         self.set_choice(choice)
         session.log_choice(self)
+        session.task_match_chosen()
 
     def reload(self):
         """Reload albums and items from the database."""
@@ -693,10 +698,11 @@ def _emit_imported(self, lib):
         for item in self.imported_items():
             plugins.send("item_imported", lib=lib, item=item)
 
-    def lookup_candidates(self, search_ids: list[str]) -> None:
+    def lookup_candidates(self, session: ImportSession, search_ids: list[str]) -> None:
         self.candidates, self.rec = autotag.tag_item(
             self.item, search_ids=search_ids
         )
+        session.task_candidates_found()
 
     def find_duplicates(self, lib: library.Library) -> list[library.Item]:  # type: ignore[override] # Need splitting Singleton and Album tasks into separate classes
         """Return a list of items from `lib` that have the same artist

beets/test/helper.py

@@ -663,13 +663,30 @@ class ImportSessionFixture(ImportSession):
     remaining albums, the metadata from the autotagger will be applied.
     """
 
+    created: int = 0
+    candidates_found: int = 0
+    match_chosen: int = 0
+    finalized: int = 0
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._choices = []
         self._resolutions = []
 
     default_choice = importer.Action.APPLY
 
+    def tasks_created(self, tasks: list[importer.ImportTask]) -> None:
+        self.created += len(tasks)
+
+    def task_candidates_found(self) -> None:
+        self.candidates_found += 1
+
+    def task_match_chosen(self) -> None:
+        self.match_chosen += 1
+
+    def task_finalized(self) -> None:
+        self.finalized += 1
+
     def add_choice(self, choice):
         self._choices.append(choice)
 
@@ -710,13 +727,30 @@ def resolve_duplicate(self, task, found_duplicates):
 
 
 class TerminalImportSessionFixture(TerminalImportSession):
+    created: int = 0
+    candidates_found: int = 0
+    match_chosen: int = 0
+    finalized: int = 0
+
     def __init__(self, *args, **kwargs):
         self.io = kwargs.pop("io")
         super().__init__(*args, **kwargs)
         self._choices = []
 
     default_choice = importer.Action.APPLY
 
+    def tasks_created(self, tasks: list[importer.ImportTask]) -> None:
+        self.created += len(tasks)
+
+    def task_candidates_found(self) -> None:
+        self.candidates_found += 1
+
+    def task_match_chosen(self) -> None:
+        self.match_chosen += 1
+
+    def task_finalized(self) -> None:
+        self.finalized += 1
+
     def add_choice(self, choice):
         self._choices.append(choice)
 

beets/ui/__init__.py

@@ -28,22 +28,26 @@
 import sys
 import textwrap
 import traceback
+from contextlib import contextmanager
 from difflib import SequenceMatcher
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any, Callable, Generator
 
 import confuse
+import enlighten
 
 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
 
+is_windows = sys.platform == "win32"
+
 if TYPE_CHECKING:
     from types import ModuleType
 
 # On Windows platforms, use colorama to support "ANSI" terminal colors.
-if sys.platform == "win32":
+if is_windows:
     try:
         import colorama
     except ImportError:
@@ -1370,6 +1374,97 @@ def add_all_common_options(self):
         self.add_format_option()
 
 
+@contextmanager
+def changes_and_errors_pbars(
+    **kwargs,
+) -> Generator[
+    tuple[enlighten.Counter, enlighten.Counter, enlighten.Counter], None, None
+]:
+    """Construct three progress bars for incremental changes and errors.
+
+    Using this method to construct the three progress bars allows Beets to
+    manage the formatting and coloring of the progress bars, ensuring
+    consistency across the codebase and among plugins.
+
+    Example usage:
+
+    ```python
+    with ui.changes_and_errors_pbars(
+        total=len(items),
+        desc="Updating items",
+        unit="items",
+    ) as (n_changed, n_unchanged, n_errors):
+        for album in lib.albums():
+            try:
+                if update_album(album):
+                    n_changed.update()
+                else:
+                    n_unchanged.update()
+            except Exception:
+                n_errors.update()
+    ```
+
+    Args:
+        kwargs: Keyword arguments to pass to the `enlighten.Counter`
+            constructor.
+
+    Returns:
+        A tuple of three `enlighten.Counter` instances: the first for changed
+        items, the second for unchanged items, and the third for errors.
+    """
+    if "color" in kwargs:
+        del kwargs["color"]
+
+    # Currently disabled when running in Windows. Enlighten does not seem to
+    # handle user inputs in the text area above the progress bar, nor cleaning
+    # up the progress bar from the terminal when the command completes.
+    #
+    # TODO: investigate and resolve these problems, and enable the progress
+    # bars in Windows environments.
+    with enlighten.Manager(enabled=not is_windows) as manager:
+        with manager.counter(**kwargs, color="white") as unchanged:
+            changed = unchanged.add_subcounter("blue")
+            errors = unchanged.add_subcounter("red")
+            yield changed, unchanged, errors
+
+
+def iprogress_bar(sequence, **kwargs):
+    """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)
+    ```
+
+    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 sequence is not None, and can return its length, then use that as the total.
+    if "total" not in kwargs and hasattr(sequence, "__len__"):
+        kwargs["total"] = len(sequence)
+
+    # Disabled in windows environments. See above for details
+    with enlighten.Manager(enabled=not is_windows) as manager:
+        with manager.counter(**kwargs) as counter:
+            for item in sequence:
+                counter.update()
+                yield item
+
+
 # Subcommand parsing infrastructure.
 #
 # This is a fairly generic subcommand parser for optparse. It is