node: Improve live progress output

This upgrades the progress output quite a bit, largely thanks to the
rich CLI library:

- There's a pretty animation at the start of the line!
- Instead of printing the last package interacted with (which made no
  sense), show as many of the currently-processing packages as can fit
  on a single line.
- Show live download progress for large downloads, so any that are
  holding up generation don't appears to be hanging.

In addition, this adds a hidden CLI flag, --traceback-on-interrupt, that
prints a full traceback from every active package coroutine on Ctrl-C,
which was useful while debugging this (and probably will be in the
future!)

Signed-off-by: Ryan Gonzalez <ryan.gonzalez@collabora.com>

node/flatpak_node_generator/main.py

@@ -1,13 +1,17 @@
 from pathlib import Path
-from typing import Iterator, List, Set
+from typing import Any, ContextManager, Iterator, List, Optional, Set
 
 import argparse
 import asyncio
+import contextlib
 import json
 import os
 import sys
 import time
 
+from rich.console import Console, Group, RenderableType
+from rich.live import Live
+
 from .cache import Cache, FilesystemBasedCache
 from .manifest import ManifestGenerator
 from .node_headers import NodeHeaders
@@ -19,6 +23,8 @@
 from .providers.yarn import YarnProviderFactory
 from .requests import Requests, StubRequests
 
+_CONSOLE_REFRESH_PER_SECOND = 12.5
+
 
 def _scan_for_lockfiles(base: Path, patterns: List[str]) -> Iterator[Path]:
     for root, _, files in os.walk(base.parent):
@@ -53,6 +59,11 @@ async def _async_main() -> None:
         action='append',
         help='Given -r, restrict files to those matching the given pattern.',
     )
+    parser.add_argument(
+        '--no-live-progress',
+        action='store_true',
+        help='Disable live progress output',
+    )
     parser.add_argument(
         '--registry',
         help='The registry to use (npm only)',
@@ -136,13 +147,20 @@ async def _async_main() -> None:
         dest='xdg_layout',
         help="Don't use the XDG layout for caches",
     )
-    # Internal option, useful for testing.
+    # Internal options, useful for testing.
     parser.add_argument('--stub-requests', action='store_true', help=argparse.SUPPRESS)
+    parser.add_argument(
+        '--traceback-on-interrupt',
+        action='store_true',
+        help=argparse.SUPPRESS,
+    )
 
     args = parser.parse_args()
 
     Requests.retries = args.retries
 
+    console = Console() if not args.no_live_progress and sys.stdout.isatty else None
+
     if args.type == 'yarn' and (args.no_devel or args.no_autopatch):
         sys.exit('--no-devel and --no-autopatch do not apply to Yarn.')
 
@@ -223,16 +241,45 @@ async def _async_main() -> None:
         )
         special = SpecialSourceProvider(gen, options)
 
-        with provider_factory.create_module_provider(gen, special) as module_provider:
-            with GeneratorProgress(
+        live: ContextManager[Any]
+        if console is not None:
+            requests_renderable = Requests.instance.get_renderable(console)
+            generator_renderable: Optional[RenderableType] = None
+
+            def get_renderable() -> RenderableType:
+                if generator_renderable is not None:
+                    return Group(
+                        requests_renderable,
+                        generator_renderable,
+                    )
+                else:
+                    return requests_renderable
+
+            live = Live(
+                get_renderable=get_renderable,
+                refresh_per_second=_CONSOLE_REFRESH_PER_SECOND,
+                console=console,
+            )
+        else:
+            live = contextlib.nullcontext()
+
+        with live:
+            with provider_factory.create_module_provider(
+                gen, special
+            ) as module_provider, GeneratorProgress(
                 packages,
                 module_provider,
-                args.max_parallel,
+                max_parallel=args.max_parallel,
+                traceback_on_interrupt=args.traceback_on_interrupt,
             ) as progress:
+                if console is not None:
+                    generator_renderable = progress.get_renderable(console)
+
                 await progress.run()
-        for headers in rcfile_node_headers:
-            print(f'Generating headers {headers.runtime} @ {headers.target}')
-            await special.generate_node_headers(headers)
+
+            for headers in rcfile_node_headers:
+                print(f'Generating headers {headers.runtime} @ {headers.target}...')
+                await special.generate_node_headers(headers)
 
         if args.xdg_layout:
             script_name = 'setup_sdk_node_headers.sh'

node/flatpak_node_generator/progress.py

@@ -1,73 +1,153 @@
-from typing import Collection, ContextManager, Optional, Type
+from dataclasses import dataclass
+from typing import Collection, ContextManager, Optional, Set, Type
 
 import asyncio
 import shutil
 import sys
+import traceback
 import types
 
+from rich.console import (
+    Console,
+    ConsoleOptions,
+    ConsoleRenderable,
+    RenderableType,
+    RenderResult,
+)
+from rich.measure import Measurement
+from rich.segment import Segment
+from rich.status import Status
+
 from .package import Package
 from .providers import ModuleProvider
 
 
+def _generating_packages(finished: int, total: int) -> str:
+    return f'Generating packages [{finished}/{total}]'
+
+
+class _GeneratingPackagesRenderable(ConsoleRenderable):
+    def __init__(self, finished: int, total: int, processing: Set[Package]) -> None:
+        self.generating_string = _generating_packages(finished, total)
+        self.processing = processing
+
+    def __rich_measure__(
+        self, console: Console, options: ConsoleOptions
+    ) -> Measurement:
+        return Measurement(0, options.max_width)
+
+    def __rich_console__(
+        self, console: Console, options: ConsoleOptions
+    ) -> RenderResult:
+        ARROW = ' => '
+        ELLIPSES = '...'
+        SEPARATOR = ', '
+
+        yield Segment(self.generating_string)
+        space_remaining = options.max_width - len(self.generating_string)
+
+        generating_string_width = len(self.generating_string)
+        if space_remaining < len(ELLIPSES):
+            return
+        elif options.max_width < len(ELLIPSES) + len(ARROW):
+            return ELLIPSES
+
+        packages = sorted(
+            f'{package.name} @ {package.version}' for package in self.processing
+        )
+
+        yield Segment(ARROW)
+        space_remaining -= len(ARROW) + len(ELLIPSES)
+
+        for i, package in enumerate(packages):
+            if i:
+                package = SEPARATOR + package
+            if len(package) > space_remaining:
+                break
+
+            yield Segment(package)
+            space_remaining -= len(package)
+
+        yield Segment(ELLIPSES)
+
+
 class GeneratorProgress(ContextManager['GeneratorProgress']):
     def __init__(
         self,
         packages: Collection[Package],
         module_provider: ModuleProvider,
+        *,
         max_parallel: int,
+        traceback_on_interrupt: bool,
     ) -> None:
         self.finished = 0
+        self.processing: Set[Package] = set()
         self.packages = packages
         self.module_provider = module_provider
         self.parallel_limit = asyncio.Semaphore(max_parallel)
-        self.previous_package: Optional[Package] = None
-        self.current_package: Optional[Package] = None
+        self.traceback_on_interrupt = traceback_on_interrupt
+        self.status: Optional[Status] = None
+
+    @property
+    def _total(self) -> int:
+        return len(self.packages)
 
     def __exit__(
         self,
         exc_type: Optional[Type[BaseException]],
         exc_value: Optional[BaseException],
         tb: Optional[types.TracebackType],
     ) -> None:
-        print()
-
-    def _format_package(self, package: Package, max_width: int) -> str:
-        result = f'{package.name} @ {package.version}'
-
-        if len(result) > max_width:
-            result = result[: max_width - 3] + '...'
-
-        return result
+        line = f'Generated {self._total} package(s).'
+        if self.status is not None:
+            self.status.update(line)
+            self.status.stop()
+        else:
+            print(line)
 
     def _update(self) -> None:
-        columns, _ = shutil.get_terminal_size()
-
-        sys.stdout.write('\r' + ' ' * columns)
-
-        prefix_string = f'\rGenerating packages [{self.finished}/{len(self.packages)}] '
-        sys.stdout.write(prefix_string)
-        max_package_width = columns - len(prefix_string)
-
-        if self.current_package is not None:
-            sys.stdout.write(
-                self._format_package(self.current_package, max_package_width)
-            )
-
-        sys.stdout.flush()
-
-    def _update_with_package(self, package: Package) -> None:
-        self.previous_package, self.current_package = (
-            self.current_package,
-            package,
+        if self.status is None:
+            # No TTY. Only print an update on multiples of 10 to avoid spamming
+            # the console.
+            if self.finished % 10 == 0 or self.finished == self._total:
+                print(
+                    f'{_generating_packages(self.finished, self._total)}...',
+                    flush=True,
+                )
+            return
+
+        self.status.update(
+            _GeneratingPackagesRenderable(self.finished, self._total, self.processing)
         )
-        self._update()
 
     async def _generate(self, package: Package) -> None:
         async with self.parallel_limit:
-            self._update_with_package(package)
-            await self.module_provider.generate_package(package)
+            self.processing.add(package)
+            # Don't bother printing an update here without live progress, since
+            # then the currently processing packages won't appear anyway.
+            if self.status is not None:
+                self._update()
+
+            try:
+                await self.module_provider.generate_package(package)
+            except asyncio.CancelledError:
+                if self.traceback_on_interrupt:
+                    print(f'========== {package.name} ==========', file=sys.stderr)
+                    traceback.print_exc()
+                    print(file=sys.stderr)
+                raise
+
             self.finished += 1
-            self._update_with_package(package)
+            self.processing.remove(package)
+            self._update()
+
+    def get_renderable(self, console: Console) -> RenderableType:
+        if self.status is not None:
+            assert self.status.console is console
+        else:
+            self.status = Status('', console=console)
+
+        return self.status
 
     async def run(self) -> None:
         self._update()