346 lines
10 KiB
Python
346 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
from collections import deque
|
|
from math import floor, log10
|
|
from threading import Lock
|
|
|
|
|
|
_WINDOWS_ANSI_ENABLED: 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 sys.stdout.isatty():
|
|
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 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() -> None:
|
|
if not sys.stdout.isatty():
|
|
return
|
|
os.system("cls" if os.name == "nt" else "clear")
|
|
|
|
|
|
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 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,
|
|
*,
|
|
stream=None,
|
|
embedded_percent: bool = False,
|
|
show_rate: bool = False,
|
|
) -> None:
|
|
self.total = max(0, total)
|
|
self.label = label
|
|
self.stream = stream or sys.stdout
|
|
self.started_at = time.monotonic()
|
|
self.samples: deque[tuple[int, float]] = deque()
|
|
self.last_width = 0
|
|
self.embedded_percent = embedded_percent
|
|
self.show_rate = show_rate
|
|
|
|
def update(self, processed: int, detail: str = "") -> None:
|
|
self._write(self.update_line(processed, detail))
|
|
|
|
def update_line(self, processed: int, detail: str = "") -> str:
|
|
processed = max(0, min(processed, self.total))
|
|
now = time.monotonic()
|
|
if not self.samples or self.samples[-1][0] != processed:
|
|
self.samples.append((processed, now))
|
|
while len(self.samples) > max(2, processed // 2 + 2):
|
|
self.samples.popleft()
|
|
return self._line(processed, detail, now)
|
|
|
|
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 self.stream.isatty():
|
|
if keep:
|
|
self.stream.write("\n")
|
|
else:
|
|
self.stream.write("\r" + " " * self.last_width + "\r")
|
|
self.stream.flush()
|
|
elif self.last_width and not keep:
|
|
self.stream.write("\n")
|
|
self.last_width = 0
|
|
|
|
def _write(self, line: str) -> None:
|
|
if self.stream.isatty():
|
|
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(line + "\n")
|
|
self.last_width = len(line)
|
|
|
|
def _line(
|
|
self,
|
|
processed: int,
|
|
detail: str,
|
|
now: float | None,
|
|
*,
|
|
elapsed: float | None = None,
|
|
eta: float | None = None,
|
|
) -> str:
|
|
prefix = self.label if not detail else f"{self.label}: {detail}"
|
|
if self.total <= 2:
|
|
return f"{prefix} {processed}/{self.total}"
|
|
|
|
percent = 0.0 if self.total == 0 else processed / self.total
|
|
bar = _progress_bar(percent, embedded_percent=self.embedded_percent)
|
|
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")
|
|
parts.append(f"elapsed {format_duration(elapsed_value)}")
|
|
eta_value = eta if eta is not None else self._eta(processed)
|
|
if processed >= 2 and processed < self.total and eta_value is not None:
|
|
parts.append(f"ETA {format_duration(eta_value)}")
|
|
return " ".join(parts)
|
|
|
|
def _eta(self, processed: int) -> float | None:
|
|
if processed < 2 or self.total <= processed or len(self.samples) < 2:
|
|
return None
|
|
oldest_processed, oldest_time = self.samples[0]
|
|
newest_processed, newest_time = self.samples[-1]
|
|
delta_items = newest_processed - oldest_processed
|
|
delta_time = newest_time - oldest_time
|
|
if delta_items <= 0 or delta_time <= 0:
|
|
return None
|
|
return (self.total - processed) * (delta_time / delta_items)
|
|
|
|
|
|
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) -> str:
|
|
width = 24
|
|
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 _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))
|