Files
media-batch-tools/tools/console.py
T

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))