495 lines
16 KiB
Python
495 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import select
|
|
import shutil
|
|
import sys
|
|
import textwrap
|
|
import time
|
|
from collections import deque
|
|
from contextlib import contextmanager
|
|
from math import floor, log10
|
|
from threading import Lock
|
|
|
|
|
|
_WINDOWS_ANSI_ENABLED: bool | None = None
|
|
_WINDOWS_CONSOLE_OUTPUT: bool | None = None
|
|
|
|
|
|
def prompt_input(prompt: str) -> str:
|
|
try:
|
|
return input(prompt)
|
|
except EOFError as exc:
|
|
raise RuntimeError("Interactive input ended before choices were complete.") from exc
|
|
|
|
|
|
def prompt_yes_no(prompt: str, default: bool) -> bool:
|
|
suffix = "Y/n" if default else "y/N"
|
|
while True:
|
|
answer = prompt_input(f"{prompt} ({suffix}) ").strip().lower()
|
|
if not answer:
|
|
return default
|
|
if answer in {"y", "yes"}:
|
|
return True
|
|
if answer in {"n", "no"}:
|
|
return False
|
|
print("Please answer y or n.")
|
|
|
|
|
|
def pause_if_interactive() -> None:
|
|
if sys.stdin.isatty():
|
|
try:
|
|
input("\nPress Enter to close...")
|
|
except EOFError:
|
|
pass
|
|
|
|
|
|
def supports_color() -> bool:
|
|
return not os.environ.get("NO_COLOR") and supports_ansi()
|
|
|
|
|
|
def supports_ansi() -> bool:
|
|
if not stream_is_interactive(sys.stdout):
|
|
return False
|
|
if os.name != "nt":
|
|
return os.environ.get("TERM", "dumb") != "dumb"
|
|
return _enable_windows_ansi() or bool(
|
|
os.environ.get("WT_SESSION")
|
|
or os.environ.get("ANSICON")
|
|
or os.environ.get("ConEmuANSI") == "ON"
|
|
or os.environ.get("TERM")
|
|
)
|
|
|
|
|
|
def stream_is_interactive(stream) -> bool:
|
|
return stream.isatty() or (os.name == "nt" and stream is sys.stdout and _windows_console_output())
|
|
|
|
|
|
def color(text: object, code: str) -> str:
|
|
value = str(text)
|
|
if not supports_color():
|
|
return value
|
|
return f"\033[{code}m{value}\033[0m"
|
|
|
|
|
|
def light_green(text: object) -> str:
|
|
return color(text, "92")
|
|
|
|
|
|
def light_blue(text: object) -> str:
|
|
return color(text, "94")
|
|
|
|
|
|
def light_red(text: object) -> str:
|
|
return color(text, "91")
|
|
|
|
|
|
def red_strikethrough(text: object) -> str:
|
|
return color(text, "91;9")
|
|
|
|
|
|
def clear_screen(*, scrollback: bool = False) -> None:
|
|
if not stream_is_interactive(sys.stdout):
|
|
return
|
|
if supports_ansi():
|
|
sys.stdout.write("\x1b[3J" if scrollback else "")
|
|
sys.stdout.write("\x1b[H\x1b[2J")
|
|
sys.stdout.flush()
|
|
return
|
|
os.system("cls" if os.name == "nt" else "clear")
|
|
|
|
|
|
def require_terminal_ui() -> None:
|
|
if not (stream_is_interactive(sys.stdin) and stream_is_interactive(sys.stdout) and supports_ansi()):
|
|
raise RuntimeError("This settings dialog requires an ANSI-capable interactive terminal.")
|
|
|
|
|
|
@contextmanager
|
|
def raw_key_input():
|
|
require_terminal_ui()
|
|
if os.name == "nt":
|
|
yield read_key
|
|
return
|
|
|
|
import termios
|
|
import tty
|
|
|
|
fd = sys.stdin.fileno()
|
|
settings = termios.tcgetattr(fd)
|
|
try:
|
|
tty.setraw(fd)
|
|
yield read_key
|
|
finally:
|
|
termios.tcsetattr(fd, termios.TCSADRAIN, settings)
|
|
|
|
|
|
def read_key() -> str:
|
|
if os.name == "nt":
|
|
import msvcrt
|
|
|
|
key = msvcrt.getwch()
|
|
if key in {"\x00", "\xe0"}:
|
|
return {"H": "up", "P": "down", "K": "left", "M": "right"}.get(msvcrt.getwch(), "")
|
|
else:
|
|
key = sys.stdin.read(1)
|
|
if key == "\x1b" and select.select([sys.stdin], [], [], 0.03)[0]:
|
|
key += sys.stdin.read(2)
|
|
return {"\x1b[A": "up", "\x1b[B": "down", "\x1b[D": "left", "\x1b[C": "right"}.get(key, "esc")
|
|
return {"\r": "enter", "\n": "enter", "\x1b": "esc", "\x08": "backspace", "\x7f": "backspace", " ": "space"}.get(key, key)
|
|
|
|
|
|
def draw_screen(lines: list[str]) -> None:
|
|
require_terminal_ui()
|
|
sys.stdout.write("\x1b[H\x1b[2J" + "\n".join(lines))
|
|
sys.stdout.flush()
|
|
|
|
|
|
def inverse(text: str) -> str:
|
|
return f"\x1b[30;47m{text}\x1b[0m"
|
|
|
|
|
|
def dim(text: str) -> str:
|
|
return f"\x1b[90m{text}\x1b[0m"
|
|
|
|
|
|
def dark_field(text: str) -> str:
|
|
return f"\x1b[30;100m{text}\x1b[0m"
|
|
|
|
|
|
def _enable_windows_ansi() -> bool:
|
|
global _WINDOWS_ANSI_ENABLED
|
|
if _WINDOWS_ANSI_ENABLED is not None:
|
|
return _WINDOWS_ANSI_ENABLED
|
|
try:
|
|
import ctypes
|
|
|
|
kernel32 = ctypes.windll.kernel32
|
|
handle = kernel32.GetStdHandle(-11)
|
|
mode = ctypes.c_uint32()
|
|
if not kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
|
|
_WINDOWS_ANSI_ENABLED = False
|
|
return False
|
|
_WINDOWS_ANSI_ENABLED = bool(kernel32.SetConsoleMode(handle, mode.value | 0x0004))
|
|
return _WINDOWS_ANSI_ENABLED
|
|
except Exception:
|
|
_WINDOWS_ANSI_ENABLED = False
|
|
return False
|
|
|
|
|
|
def _windows_console_output() -> bool:
|
|
global _WINDOWS_CONSOLE_OUTPUT
|
|
if _WINDOWS_CONSOLE_OUTPUT is not None:
|
|
return _WINDOWS_CONSOLE_OUTPUT
|
|
try:
|
|
import ctypes
|
|
|
|
mode = ctypes.c_uint32()
|
|
handle = ctypes.windll.kernel32.GetStdHandle(-11)
|
|
_WINDOWS_CONSOLE_OUTPUT = bool(
|
|
ctypes.windll.kernel32.GetConsoleMode(handle, ctypes.byref(mode))
|
|
)
|
|
except Exception:
|
|
_WINDOWS_CONSOLE_OUTPUT = False
|
|
return _WINDOWS_CONSOLE_OUTPUT
|
|
|
|
|
|
def format_duration(seconds: float) -> str:
|
|
seconds = max(0, int(round(seconds)))
|
|
minutes, second = divmod(seconds, 60)
|
|
hours, minute = divmod(minutes, 60)
|
|
if hours:
|
|
return f"{hours}:{minute:02d}:{second:02d}"
|
|
return f"{minute}:{second:02d}"
|
|
|
|
|
|
class ProgressView:
|
|
def __init__(
|
|
self,
|
|
total: int,
|
|
label: str,
|
|
*,
|
|
total_work: float | None = None,
|
|
stream=None,
|
|
embedded_percent: bool = False,
|
|
show_rate: bool = False,
|
|
show_elapsed: bool = True,
|
|
bar_width: int = 24,
|
|
) -> None:
|
|
self.total = max(0, total)
|
|
self.total_work = total if total_work is None else max(0.0, total_work)
|
|
self.label = label
|
|
self.stream = stream or sys.stdout
|
|
self.started_at = time.monotonic()
|
|
self.samples: deque[tuple[float, float]] = deque()
|
|
self.last_width = 0
|
|
self.last_line_count = 0
|
|
self.embedded_percent = embedded_percent
|
|
self.show_rate = show_rate
|
|
self.show_elapsed = show_elapsed
|
|
self.bar_width = bar_width
|
|
|
|
def update(
|
|
self, processed: int, detail: str = "", *, completed_work: float | None = None
|
|
) -> None:
|
|
self._write(self.update_line(processed, detail, completed_work=completed_work))
|
|
|
|
def update_line(
|
|
self, processed: int, detail: str = "", *, completed_work: float | None = None
|
|
) -> str:
|
|
processed = max(0, min(processed, self.total))
|
|
work_done = min(self.total_work, completed_work if completed_work is not None else processed)
|
|
now = time.monotonic()
|
|
if not self.samples or self.samples[-1][0] != work_done:
|
|
self.samples.append((work_done, now))
|
|
while len(self.samples) > max(2, processed // 2 + 2):
|
|
self.samples.popleft()
|
|
return self._line(processed, detail, now, completed_work=work_done)
|
|
|
|
def update_external(
|
|
self,
|
|
processed: int,
|
|
*,
|
|
elapsed: float,
|
|
eta: float | None,
|
|
detail: str = "",
|
|
) -> None:
|
|
self._write(self.update_external_line(processed, elapsed=elapsed, eta=eta, detail=detail))
|
|
|
|
def update_external_line(
|
|
self,
|
|
processed: int,
|
|
*,
|
|
elapsed: float,
|
|
eta: float | None,
|
|
detail: str = "",
|
|
) -> str:
|
|
processed = max(0, min(processed, self.total))
|
|
return self._line(processed, detail, None, elapsed=elapsed, eta=eta)
|
|
|
|
def finish(self, *, keep: bool = False) -> None:
|
|
if stream_is_interactive(self.stream):
|
|
if keep:
|
|
self.stream.write("\n")
|
|
else:
|
|
self._clear_live_lines()
|
|
self.stream.flush()
|
|
elif self.last_width and not keep:
|
|
self.stream.write("\n")
|
|
self.last_width = 0
|
|
self.last_line_count = 0
|
|
|
|
def write_lines(self, lines: list[str]) -> None:
|
|
self._write(lines)
|
|
|
|
def _write(self, content: str | list[str]) -> None:
|
|
lines = [content] if isinstance(content, str) else content
|
|
if stream_is_interactive(self.stream) and supports_ansi():
|
|
lines = [wrapped for line in lines for wrapped in _wrap_to_terminal(line, self.stream)]
|
|
if self.last_line_count > 1:
|
|
self.stream.write(f"\x1b[{self.last_line_count - 1}F")
|
|
rows = max(self.last_line_count, len(lines))
|
|
for index in range(rows):
|
|
self.stream.write("\r\x1b[2K")
|
|
if index < len(lines):
|
|
self.stream.write(lines[index])
|
|
if index < rows - 1:
|
|
self.stream.write("\n")
|
|
if rows > len(lines):
|
|
self.stream.write(f"\x1b[{rows - len(lines)}F")
|
|
self.stream.flush()
|
|
self.last_width = max(map(len, lines), default=0)
|
|
self.last_line_count = len(lines)
|
|
elif stream_is_interactive(self.stream):
|
|
line = " ".join(lines)
|
|
padding = max(0, self.last_width - len(line))
|
|
self.stream.write("\r" + line + (" " * padding))
|
|
self.stream.flush()
|
|
self.last_width = len(line)
|
|
else:
|
|
self.stream.write("\n".join(lines) + "\n")
|
|
self.last_width = max(map(len, lines), default=0)
|
|
|
|
def _clear_live_lines(self) -> None:
|
|
if not self.last_line_count or not supports_ansi():
|
|
self.stream.write("\r" + " " * self.last_width + "\r")
|
|
return
|
|
if self.last_line_count > 1:
|
|
self.stream.write(f"\x1b[{self.last_line_count - 1}F")
|
|
for index in range(self.last_line_count):
|
|
self.stream.write("\r\x1b[2K")
|
|
if index < self.last_line_count - 1:
|
|
self.stream.write("\n")
|
|
if self.last_line_count > 1:
|
|
self.stream.write(f"\x1b[{self.last_line_count - 1}F")
|
|
|
|
def _line(
|
|
self,
|
|
processed: int,
|
|
detail: str,
|
|
now: float | None,
|
|
*,
|
|
elapsed: float | None = None,
|
|
eta: float | None = None,
|
|
completed_work: float | None = None,
|
|
) -> str:
|
|
prefix = self.label
|
|
if self.total <= 2:
|
|
return " ".join(part for part in (prefix, f"{processed}/{self.total}", detail) if part)
|
|
|
|
work_done = processed if completed_work is None else completed_work
|
|
percent = 0.0 if self.total_work == 0 else work_done / self.total_work
|
|
bar = _progress_bar(percent, embedded_percent=self.embedded_percent, width=self.bar_width)
|
|
elapsed_value = elapsed if elapsed is not None else (now or time.monotonic()) - self.started_at
|
|
parts = [
|
|
prefix,
|
|
bar,
|
|
f"{processed}/{self.total}",
|
|
]
|
|
if not self.embedded_percent:
|
|
parts.append(f"{percent * 100:5.1f}%")
|
|
if self.show_rate and processed > 0 and elapsed_value > 0:
|
|
parts.append(f"{_format_significant(processed / elapsed_value, 2)} fps")
|
|
if self.show_elapsed:
|
|
parts.append(f"elapsed {format_duration(elapsed_value)}")
|
|
eta_value = eta if eta is not None else self._eta(processed, work_done)
|
|
if processed >= 2 and processed < self.total and eta_value is not None:
|
|
parts.append(f"ETA {format_duration(eta_value)}")
|
|
if detail:
|
|
parts.append(detail)
|
|
return " ".join(parts)
|
|
|
|
def _eta(self, processed: int, work_done: float) -> float | None:
|
|
if processed < 2 or self.total <= processed or len(self.samples) < 2:
|
|
return None
|
|
oldest_work, oldest_time = self.samples[0]
|
|
newest_work, newest_time = self.samples[-1]
|
|
delta_work = newest_work - oldest_work
|
|
delta_time = newest_time - oldest_time
|
|
if delta_work <= 0 or delta_time <= 0:
|
|
return None
|
|
return (self.total_work - work_done) * (delta_time / delta_work)
|
|
|
|
|
|
class PipelineProgressView:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
rendering_label: str,
|
|
encoding_label: str,
|
|
encoding_total: int,
|
|
stream=None,
|
|
) -> None:
|
|
self.stream = stream or sys.stdout
|
|
self.rendering_label = rendering_label
|
|
self.encoding = ProgressView(
|
|
encoding_total,
|
|
encoding_label,
|
|
stream=self.stream,
|
|
embedded_percent=True,
|
|
show_rate=True,
|
|
)
|
|
self.rendering: ProgressView | None = None
|
|
self.rendering_line = f"{rendering_label} starting..."
|
|
self.encoding_line = self.encoding.update_line(0)
|
|
self._shown = False
|
|
self._last_width = 0
|
|
self._ansi = supports_ansi()
|
|
self._lock = Lock()
|
|
|
|
def update_rendering(
|
|
self,
|
|
processed: int,
|
|
total: int,
|
|
*,
|
|
elapsed: float,
|
|
eta: float | None,
|
|
) -> None:
|
|
with self._lock:
|
|
if self.rendering is None or self.rendering.total != total:
|
|
self.rendering = ProgressView(
|
|
total,
|
|
self.rendering_label,
|
|
stream=self.stream,
|
|
embedded_percent=True,
|
|
show_rate=True,
|
|
)
|
|
self.rendering_line = self.rendering.update_external_line(
|
|
processed,
|
|
elapsed=elapsed,
|
|
eta=eta if processed >= 2 else None,
|
|
)
|
|
self._write_locked()
|
|
|
|
def update_encoding(self, processed: int) -> None:
|
|
with self._lock:
|
|
self.encoding_line = self.encoding.update_line(processed)
|
|
self._write_locked()
|
|
|
|
def finish(self) -> None:
|
|
with self._lock:
|
|
if not self._shown:
|
|
return
|
|
if self._ansi:
|
|
self.stream.write("\n")
|
|
else:
|
|
self.stream.write("\r" + " " * self._last_width + "\r")
|
|
self.stream.flush()
|
|
self._shown = False
|
|
|
|
def _write_locked(self) -> None:
|
|
if self._ansi:
|
|
if self._shown:
|
|
self.stream.write("\x1b[1F")
|
|
self.stream.write("\r\x1b[2K" + self.rendering_line + "\n")
|
|
self.stream.write("\r\x1b[2K" + self.encoding_line)
|
|
else:
|
|
line = f"{self.rendering_line} | {self.encoding_line}"
|
|
padding = max(0, self._last_width - len(line))
|
|
self.stream.write("\r" + line + " " * padding)
|
|
self._last_width = len(line)
|
|
self.stream.flush()
|
|
self._shown = True
|
|
|
|
|
|
def _progress_bar(percent: float, *, embedded_percent: bool, width: int = 24) -> str:
|
|
done = round(max(0.0, min(percent, 1.0)) * width)
|
|
if not embedded_percent:
|
|
return "[" + ("#" * done).ljust(width, "-") + "]"
|
|
|
|
remaining = width - done
|
|
percent_text = f"{percent * 100:.1f}%"
|
|
if percent >= 0.5:
|
|
left = _center_progress_text(percent_text, done, "#")
|
|
right = "-" * remaining
|
|
else:
|
|
left = "#" * done
|
|
right = _center_progress_text(percent_text, remaining, "-")
|
|
return f"[{left}{right}]"
|
|
|
|
|
|
def _center_progress_text(text: str, width: int, fill: str) -> str:
|
|
content = f" {text} "
|
|
if len(content) > width:
|
|
return fill * width
|
|
return content.center(width, fill)
|
|
|
|
|
|
def _wrap_to_terminal(line: str, stream) -> list[str]:
|
|
try:
|
|
width = shutil.get_terminal_size().columns
|
|
except OSError:
|
|
return [line]
|
|
if width <= 0 or len(line) <= width:
|
|
return [line]
|
|
return textwrap.wrap(line, width=width, break_long_words=False, break_on_hyphens=False) or [line]
|
|
|
|
|
|
def _format_significant(value: float, digits: int) -> str:
|
|
if value == 0:
|
|
return "0"
|
|
decimals = digits - floor(log10(abs(value))) - 1
|
|
if decimals > 0:
|
|
return f"{value:.{decimals}f}"
|
|
factor = 10 ** -decimals
|
|
return str(int(round(value / factor) * factor))
|