diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 650e3b6d..75908004 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -29,10 +29,9 @@ from click.core import ParameterSource from langcodes import Language from pymediainfo import MediaInfo from rich.console import Group -from rich.live import Live from rich.padding import Padding from rich.panel import Panel -from rich.progress import BarColumn, Progress, SpinnerColumn, TaskID, TextColumn, TimeRemainingColumn +from rich.progress import Progress, SpinnerColumn, TaskID, TextColumn, TimeRemainingColumn from rich.rule import Rule from rich.spinner import Spinner from rich.table import Table @@ -43,7 +42,7 @@ from unshackle.core import binaries, providers from unshackle.core.cdm import DecryptLabsRemoteCDM from unshackle.core.cdm.detect import is_playready_cdm, is_widevine_cdm from unshackle.core.config import config, resolve_cdm_name, resolve_decryption -from unshackle.core.console import console +from unshackle.core.console import GradientPulseBarColumn, SyncLive, console from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack, context_settings from unshackle.core.credential import Credential from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine @@ -1680,7 +1679,7 @@ class dl: track = song.tracks.audio[0] progress = Progress( SpinnerColumn(finished_text=""), - BarColumn(), + GradientPulseBarColumn(), " | ", TimeRemainingColumn(compact=True, elapsed_when_finished=True), " | ", @@ -1692,17 +1691,25 @@ class dl: state = {"total": 100.0} def update_track_progress( - task_id: int = task, + task_id: TaskID = task, _state: dict[str, float] = state, _progress: Progress = progress, - **kwargs, + **kwargs: Any, ) -> None: - if "total" in kwargs and kwargs["total"] is not None: - _state["total"] = kwargs["total"] + if "total" in kwargs: + if kwargs["total"] is None: + # Progress.update() ignores total=None; an un-started task pulses + del kwargs["total"] + _progress.reset(task_id, start=False) + else: + _state["total"] = kwargs["total"] + _progress.start_task(task_id) downloaded_state = kwargs.get("downloaded") if downloaded_state in {"Downloaded", "Decrypted", "[yellow]SKIPPED"}: kwargs["completed"] = _state["total"] + kwargs["total"] = _state["total"] + _progress.start_task(task_id) _progress.update(task_id=task_id, **kwargs) track_table = Table.grid() @@ -1720,7 +1727,7 @@ class dl: download_table.add_row(music_tree) try: - with Live(Padding(download_table, (1, 5)), console=console, refresh_per_second=5): + with SyncLive(Padding(download_table, (1, 5)), console=console, refresh_per_second=20): with ThreadPoolExecutor(downloads) as pool: download_futures = [ pool.submit( @@ -2065,7 +2072,7 @@ class dl: if slow and i != 0: delay = random.randint(slow[0], slow[1]) spinner = Spinner("dots", text=f"Delaying by {delay} seconds...") - with Live(Padding(spinner, (0, 5)), console=console, refresh_per_second=12.5, transient=True): + with SyncLive(Padding(spinner, (0, 5)), console=console, refresh_per_second=12.5, transient=True): for remaining in range(delay, 0, -1): spinner.update(text=f"Delaying by {remaining} seconds...") time.sleep(1) @@ -2825,7 +2832,7 @@ class dl: dl_start_time = time.time() try: - with Live(Padding(download_table, (1, 5)), console=console, refresh_per_second=5): + with SyncLive(Padding(download_table, (1, 5)), console=console, refresh_per_second=20): def download_track(track: AnyTrack, i: int) -> None: track.download( @@ -3104,7 +3111,7 @@ class dl: progress = Progress( TextColumn("[progress.description]{task.description}"), SpinnerColumn(finished_text=""), - BarColumn(), + GradientPulseBarColumn(), "•", TimeRemainingColumn(compact=True, elapsed_when_finished=True), "•", @@ -3252,10 +3259,9 @@ class dl: {"phase": "muxing", "progress": 96.0, "status": "downloading", "active_tracks": []} ) try: - with Live(Padding(progress, (0, 5, 1, 5)), console=console): + with SyncLive(Padding(progress, (0, 5, 1, 5)), console=console, refresh_per_second=20): mux_index = 0 for task_id, task_tracks, audio_codec in multiplex_tasks: - progress.start_task(task_id) # TODO: Needed? audio_expected = not video_only and not no_audio muxed_path, return_code, errors = task_tracks.mux( str(title), @@ -3265,6 +3271,8 @@ class dl: title_language=title.language, skip_subtitles=skip_subtitle_mux, ) + progress.start_task(task_id) + progress.update(task_id, total=100, completed=100) if muxed_path.exists(): mux_index += 1 unique_path = muxed_path.with_name( diff --git a/unshackle/core/console.py b/unshackle/core/console.py index c5af3686..52c0415f 100644 --- a/unshackle/core/console.py +++ b/unshackle/core/console.py @@ -1,17 +1,21 @@ import logging +import math from datetime import datetime from types import ModuleType -from typing import IO, Callable, Iterable, List, Literal, Mapping, Optional, TextIO, Union +from typing import IO, Any, Callable, Iterable, List, Literal, Mapping, Optional, TextIO, Union from rich._log_render import FormatTimeCallable, LogRender -from rich.console import Console, ConsoleRenderable, HighlighterType, RenderableType +from rich.color import Color, blend_rgb +from rich.console import Console, ConsoleOptions, ConsoleRenderable, HighlighterType, RenderableType, RenderResult from rich.emoji import EmojiVariant from rich.highlighter import Highlighter, ReprHighlighter from rich.live import Live from rich.logging import RichHandler +from rich.measure import Measurement from rich.padding import Padding, PaddingDimensions +from rich.progress import BarColumn, Task from rich.status import Status -from rich.style import StyleType +from rich.style import Style, StyleType from rich.table import Table from rich.text import Text, TextType from rich.theme import Theme @@ -279,7 +283,7 @@ class ComfyConsole(Console): padding = Padding(status_renderable, (top, right, bottom, left)) - return Live(padding, console=self, transient=True) + return SyncLive(padding, console=self, transient=True) return status_renderable @@ -349,6 +353,60 @@ if config.set_terminal_bg: custom_colors["ascii.art"] += f" on {primary_scheme['bg']}" +class SyncLive(Live): + """Live that wraps each repaint in DEC 2026 begin/end-sync so the terminal draws it as one atomic frame.""" + + def refresh(self) -> None: + if not self.console.is_terminal: + super().refresh() + return + self.console.file.write("\x1b[?2026h") + try: + super().refresh() + finally: + self.console.file.write("\x1b[?2026l") + self.console.file.flush() + + +class GradientPulseBarColumn(BarColumn): + _LUT_SIZE = 32 + + def __init__(self, *args: Any, pulse_period: float = 40.0, pulse_speed: float = 15.0, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.pulse_period = pulse_period + self.pulse_speed = pulse_speed + lo = Color.parse(primary_scheme["dark_gray"]).triplet + hi = Color.parse(primary_scheme["pink"]).triplet + assert lo and hi + self._lut = [ + Style(color=Color.from_triplet(blend_rgb(lo, hi, i / (self._LUT_SIZE - 1)))) for i in range(self._LUT_SIZE) + ] + + def render(self, task: Task) -> RenderableType: # type: ignore[override] + if task.started and task.total is not None: + return super().render(task) + return _GradientPulse(self._lut, task.get_time() * self.pulse_speed, self.pulse_period, self.bar_width) + + +class _GradientPulse: + def __init__(self, lut: list[Style], offset: float, period: float, width: Optional[int]) -> None: + self.lut = lut + self.offset = offset + self.period = period + self.width = width + + def __rich_console__(self, console: Console, options: ConsoleOptions) -> RenderResult: + width = min(self.width or options.max_width, options.max_width) + bar = Text() + for i in range(width): + fade = (math.cos((i - self.offset) * math.tau / self.period) + 1) / 2 + bar.append("━", self.lut[int(fade * (len(self.lut) - 1))]) + yield bar + + def __rich_measure__(self, console: Console, options: ConsoleOptions) -> Measurement: + return Measurement(4, self.width or options.max_width) + + console = ComfyConsole( log_time=False, log_path=False, @@ -358,7 +416,7 @@ console = ComfyConsole( "bar.back": primary_scheme["dark_gray"], "bar.complete": primary_scheme["pink"], "bar.finished": primary_scheme["green"], - "bar.pulse": primary_scheme["bright_black"], + "bar.pulse": primary_scheme["pink"], "black": primary_scheme["black"], "inspect.async_def": f"italic {primary_scheme['cyan']}", "progress.data.speed": "dark_orange", @@ -377,4 +435,4 @@ console = ComfyConsole( ) -__all__ = ("ComfyLogRenderer", "ComfyRichHandler", "ComfyConsole", "console") +__all__ = ("ComfyLogRenderer", "ComfyRichHandler", "ComfyConsole", "GradientPulseBarColumn", "SyncLive", "console") diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index 1d15cf5b..66f28374 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -547,7 +547,7 @@ class DASH: if init_data: f.write(init_data) if len(segments_to_merge) > 1: - progress(downloaded="Merging", completed=0, total=len(segments_to_merge)) + progress(downloaded="Merging", completed=0, total=None) for segment_file in segments_to_merge: if is_text_subtitle: segment_data = try_ensure_utf8(segment_file.read_bytes()) @@ -568,11 +568,11 @@ class DASH: events.emit(events.Types.TRACK_DOWNLOADED, track=track) if drm: - progress(downloaded="Decrypting", completed=0, total=100) + progress(downloaded="Decrypting", completed=0, total=None) drm.decrypt(save_path) track.drm = None events.emit(events.Types.TRACK_DECRYPTED, track=track, drm=drm, segment=None) - progress(downloaded="Decrypting", advance=100) + progress(downloaded="Decrypted", completed=100, total=100) # Clean up empty segment directory if save_dir.exists() and save_dir.name.endswith("_segments"): diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index 2b1c2eec..91933a48 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -701,7 +701,7 @@ class HLS: for control_file in segment_save_dir.glob("*.!dev"): control_file.unlink(missing_ok=True) - progress(total=total_segments, completed=0, downloaded="Merging") + progress(downloaded="Merging", completed=0, total=None) name_len = len(str(total_segments)) discon_i = 0 @@ -999,7 +999,7 @@ class HLS: if len(segments_to_merge) == 1: shutil.move(segments_to_merge[0], save_path) else: - progress(downloaded="Merging") + progress(downloaded="Merging", completed=0, total=None) if isinstance(track, (Video, Audio)): HLS.merge_segments(segments=segments_to_merge, save_path=save_path) else: @@ -1019,7 +1019,7 @@ class HLS: # Directory might not be empty, try removing recursively shutil.rmtree(save_dir, ignore_errors=True) - progress(downloaded="Downloaded") + progress(downloaded="Downloaded", completed=100, total=100) track.path = save_path diff --git a/unshackle/core/manifests/ism.py b/unshackle/core/manifests/ism.py index 0c8913ef..21b98222 100644 --- a/unshackle/core/manifests/ism.py +++ b/unshackle/core/manifests/ism.py @@ -497,6 +497,7 @@ class ISM: and isinstance(track, Subtitle) and track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML) ) + progress(downloaded="Merging", completed=0, total=None) with open(save_path, "wb") as f: first_segment = segments_to_merge[0].read_bytes() if segments_to_merge else None init_segment = ISM._init_segment(track, session_drm, first_segment) @@ -526,11 +527,11 @@ class ISM: events.emit(events.Types.TRACK_DOWNLOADED, track=track) if session_drm: - progress(downloaded="Decrypting", completed=0, total=100) + progress(downloaded="Decrypting", completed=0, total=None) session_drm.decrypt(save_path) track.drm = None events.emit(events.Types.TRACK_DECRYPTED, track=track, drm=session_drm, segment=None) - progress(downloaded="Decrypting", advance=100) + progress(downloaded="Decrypted", completed=100, total=100) try: save_dir.rmdir() diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index 11e09c78..f7c384cb 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -374,11 +374,11 @@ class Track: events.emit(events.Types.TRACK_DOWNLOADED, track=self) if drm: - progress(downloaded="Decrypting", completed=0, total=100) + progress(downloaded="Decrypting", completed=0, total=None) drm.decrypt(save_path) self.drm = None events.emit(events.Types.TRACK_DECRYPTED, track=self, drm=drm, segment=None) - progress(downloaded="Decrypted", completed=100) + progress(downloaded="Decrypted", completed=100, total=100) if track_type == "Subtitle" and self.codec.name not in ("fVTT", "fTTML"): track_data = self.path.read_bytes() diff --git a/unshackle/core/tracks/tracks.py b/unshackle/core/tracks/tracks.py index 88d04cf7..8fdcde3e 100644 --- a/unshackle/core/tracks/tracks.py +++ b/unshackle/core/tracks/tracks.py @@ -5,16 +5,16 @@ import subprocess import time from functools import partial from pathlib import Path -from typing import Callable, Iterator, Optional, Sequence, Union +from typing import Any, Callable, Iterator, Optional, Sequence, Union from langcodes import Language, closest_supported_match -from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn, TimeRemainingColumn +from rich.progress import Progress, SpinnerColumn, TaskID, TextColumn, TimeRemainingColumn from rich.table import Table from rich.tree import Tree from unshackle.core import binaries from unshackle.core.config import config -from unshackle.core.console import console +from unshackle.core.console import GradientPulseBarColumn, console from unshackle.core.constants import LANGUAGE_EXACT_DISTANCE, LANGUAGE_MAX_DISTANCE, AnyTrack, TrackT from unshackle.core.events import events from unshackle.core.tracks.attachment import Attachment @@ -114,7 +114,7 @@ class Tracks: if add_progress and track_type not in (Chapter, Attachment): progress = Progress( SpinnerColumn(finished_text=""), - BarColumn(), + GradientPulseBarColumn(), "•", TimeRemainingColumn(compact=True, elapsed_when_finished=True), "•", @@ -126,10 +126,10 @@ class Tracks: state = {"total": 100.0} def update_track_progress( - task_id: int = task, + task_id: TaskID = task, _state: dict[str, float] = state, _progress: Progress = progress, - **kwargs, + **kwargs: Any, ) -> None: """ Ensure terminal status states render as a fully completed bar. @@ -137,12 +137,20 @@ class Tracks: Some downloaders can report completed slightly below total before emitting the final "Downloaded" state. """ - if "total" in kwargs and kwargs["total"] is not None: - _state["total"] = kwargs["total"] + if "total" in kwargs: + if kwargs["total"] is None: + # Progress.update() ignores total=None; an un-started task pulses + del kwargs["total"] + _progress.reset(task_id, start=False) + else: + _state["total"] = kwargs["total"] + _progress.start_task(task_id) downloaded_state = kwargs.get("downloaded") if downloaded_state in {"Downloaded", "Decrypted", "[yellow]SKIPPED"}: kwargs["completed"] = _state["total"] + kwargs["total"] = _state["total"] + _progress.start_task(task_id) _progress.update(task_id=task_id, **kwargs) progress_callables.append(update_track_progress)