xref: /aosp_15_r20/external/crosvm/tools/impl/tui.py (revision bb4ee6a4ae7042d18b07a98463b9c8b875e44b39)
1#!/usr/bin/env python3
2# Copyright 2023 The ChromiumOS Authors
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""
7Implements styles for `Command.fg(style=)` that use `rich` terminal UI features.
8"""
9
10import subprocess
11from typing import List
12
13from .util import ensure_packages_exist
14
15ensure_packages_exist("rich")
16import rich
17import rich.console
18import rich.live
19import rich.spinner
20import rich.text
21
22
23class Styles(object):
24    "A collection of methods that can be passed to `Command.fg(style=)`"
25
26    @staticmethod
27    def live_truncated(num_lines: int = 8):
28        "Prints only the last `num_lines` of output while the program is running and successful."
29
30        def output(process: "subprocess.Popen[str]"):
31            assert process.stdout
32            spinner = rich.spinner.Spinner("dots")
33            lines: List[rich.text.Text] = []
34            stdout: List[str] = []
35            with rich.live.Live(refresh_per_second=30, transient=True) as live:
36                for line in iter(process.stdout.readline, ""):
37                    stdout.append(line.strip())
38                    lines.append(rich.text.Text.from_ansi(line.strip(), no_wrap=True))
39                    while len(lines) > num_lines:
40                        lines.pop(0)
41                    live.update(rich.console.Group(rich.text.Text("…"), *lines, spinner))
42            if process.wait() == 0:
43                console.print(rich.console.Group(rich.text.Text("…"), *lines))
44            else:
45                for line in stdout:
46                    print(line)
47
48        return output
49
50    @staticmethod
51    def quiet_with_progress(title: str):
52        "Prints only the last `num_lines` of output while the program is running and successful."
53
54        def output(process: "subprocess.Popen[str]"):
55            assert process.stdout
56            with rich.live.Live(
57                rich.spinner.Spinner("dots", title), refresh_per_second=30, transient=True
58            ):
59                stdout = process.stdout.read()
60
61            if process.wait() == 0:
62                console.print(f"[green]OK[/green] {title}")
63            else:
64                print(stdout)
65                console.print(f"[red]ERR[/red] {title}")
66
67        return output
68
69
70console = rich.console.Console()
71