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

@@ -248,6 +248,35 @@ def _setup_logging(self, loghandler: logging.Handler | None):
         logger.handlers = [loghandler]
         return logger
 
+    def tasks_created(self, _: list[ImportTask]) -> None:
+        """Called when a list of tasks is created.
+
+        Expected to be called when an individual directory or query result is
+        transformed into a list of tasks.
+        """
+        raise NotImplementedError
+
+    def task_candidates_found(self) -> None:
+        """Called when a task has found candidates.
+
+        Expected to be called by an ImportTask when it has found candidates.
+        """
+        raise NotImplementedError
+
+    def task_match_chosen(self) -> None:
+        """Called when a task has chosen a match.
+
+        Expected to be called by an ImportTask when it has chosen a match.
+        """
+        raise NotImplementedError
+
+    def task_finalized(self) -> None:
+        """Called when a task has been finalized.
+
+        Expected to be called by an ImportTask when it has been finalized.
+        """
+        raise NotImplementedError
+
     def set_config(self, config):
         """Set `config` property from global import config and make
         implied changes.
@@ -610,7 +639,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.
@@ -625,7 +654,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."""
@@ -694,6 +723,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.
@@ -731,9 +762,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):
+    def lookup_candidates(self, session: ImportSession):
         """Retrieve and store candidates for this album. User-specified
         candidate IDs are stored in self.search_ids: if present, the
         initial lookup is restricted to only those IDs.
@@ -745,6 +777,7 @@ def lookup_candidates(self):
         self.cur_album = album
         self.candidates = prop.candidates
         self.rec = prop.recommendation
+        session.task_candidates_found()
 
     def find_duplicates(self, lib: library.Library):
         """Return a list of albums from `lib` with the same artist and
@@ -1017,6 +1050,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."""
@@ -1072,10 +1106,11 @@ def _emit_imported(self, lib):
         for item in self.imported_items():
             plugins.send("item_imported", lib=lib, item=item)
 
-    def lookup_candidates(self):
+    def lookup_candidates(self, session: ImportSession):
         prop = autotag.tag_item(self.item, search_ids=self.search_ids)
         self.candidates = prop.candidates
         self.rec = prop.recommendation
+        session.task_candidates_found()
 
     def find_duplicates(self, lib):
         """Return a list of items from `lib` that have the same artist
@@ -1516,7 +1551,7 @@ def read_tasks(session: ImportSession):
         log.info("Skipped {0} paths.", skipped)
 
 
-def query_tasks(session: ImportSession):
+def query_tasks(session: ImportSession) -> Iterable[ImportTask]:
     """A generator that works as a drop-in-replacement for read_tasks.
     Instead of finding files from the filesystem, a query is used to
     match items from the library.
@@ -1564,7 +1599,7 @@ def lookup_candidates(session: ImportSession, task: ImportTask):
     # option. Currently all the IDs are passed onto the tasks directly.
     task.search_ids = session.config["search_ids"].as_str_seq()
 
-    task.lookup_candidates()
+    task.lookup_candidates(session)
 
 
 @pipeline.stage

beets/test/helper.py

@@ -679,13 +679,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)
 
@@ -726,13 +743,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

@@ -26,10 +26,12 @@
 import sys
 import textwrap
 import traceback
+from contextlib import contextmanager
 from difflib import SequenceMatcher
-from typing import Any, Callable
+from typing import Any, Callable, Generator
 
 import confuse
+import enlighten
 
 from beets import config, library, logging, plugins, util
 from beets.autotag import mb
@@ -38,8 +40,10 @@
 from beets.util import as_string
 from beets.util.functemplate import template
 
+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:
@@ -1436,6 +1440,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