From 094c598ef5d67eb8ff1477d34b59abc8b5384562 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Thu, 16 Jul 2026 11:02:34 -0600 Subject: [PATCH] perf: little improvements here and there --- pyproject.toml | 5 +++ unshackle/commands/dl.py | 8 +++- unshackle/core/binaries.py | 62 ++++++++++++++++---------- unshackle/core/downloaders/requests.py | 20 +++++++-- unshackle/core/manifests/dash.py | 12 ++++- unshackle/core/manifests/hls.py | 2 +- unshackle/core/tracks/track.py | 30 ++++++++++--- unshackle/core/tracks/video.py | 44 +++++++++++------- 8 files changed, 131 insertions(+), 52 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bdaefc53..a78d0fc4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -141,6 +141,11 @@ line-length = 120 [tool.ruff.lint] select = ["E4", "E7", "E9", "F", "W"] +[tool.ruff.lint.per-file-ignores] +# binaries are resolved lazily via module-level __getattr__ (PEP 562), so the +# names in __all__ are intentionally not defined statically. +"unshackle/core/binaries.py" = ["F822"] + [tool.isort] profile = "black" line_length = 118 diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 88c49e7e..a8c71551 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -3141,9 +3141,13 @@ class dl: ) with console.status("Repackaging tracks with FFMPEG..."): has_repacked = False + # Video tracks whose VUI was folded into the repack pass; skipped in the VUI loop below. + vui_folded = set() for track in title.tracks: if track.needs_repack: - track.repackage() + bsf = track.vui_bsf() if isinstance(track, Video) else None + if track.repackage(bsf_v=bsf): + vui_folded.add(id(track)) has_repacked = True events.emit(events.Types.TRACK_REPACKED, track=track) if has_repacked: @@ -3152,6 +3156,8 @@ class dl: with console.status("Normalizing video VUI..."): for track in title.tracks.videos: + if id(track) in vui_folded: + continue try: track.normalize_vui() except Exception as e: # noqa: BLE001 diff --git a/unshackle/core/binaries.py b/unshackle/core/binaries.py index 5ddf3a02..daa77506 100644 --- a/unshackle/core/binaries.py +++ b/unshackle/core/binaries.py @@ -34,30 +34,44 @@ def find(*names: str) -> Optional[Path]: return None -FFMPEG = find("ffmpeg") -FFProbe = find("ffprobe") -FFPlay = find("ffplay") -# seconv = SubtitleEdit 5+ CLI (the 5.0 GUI has no batch mode); SubtitleEdit = 4.x -SubtitleEdit = find("seconv", "SubtitleEdit") -ShakaPackager = find( - "shaka-packager", - "packager", - f"packager-{__shaka_platform}", - f"packager-{__shaka_platform}-arm64", - f"packager-{__shaka_platform}-x64", -) -CCExtractor = find("ccextractor", "ccextractorwin", "ccextractorwinfull") -HolaProxy = find("hola-proxy") -MPV = find("mpv") -Caddy = find("caddy") -MKVToolNix = find("mkvmerge") -Mkvpropedit = find("mkvpropedit") -DoviTool = find("dovi_tool") -HDR10PlusTool = find("hdr10plus_tool", "HDR10Plus_tool") -Mp4decrypt = find("mp4decrypt") -Docker = find("docker") -ML_Worker = find("ML-Worker") -Git = find("git") +# Binary attribute -> candidate names passed to find(). Resolved lazily on first +# attribute access (see __getattr__) to avoid ~500ms of shutil.which PATH scans at import. +__binaries = { + "FFMPEG": ("ffmpeg",), + "FFProbe": ("ffprobe",), + "FFPlay": ("ffplay",), + # seconv = SubtitleEdit 5+ CLI (the 5.0 GUI has no batch mode); SubtitleEdit = 4.x + "SubtitleEdit": ("seconv", "SubtitleEdit"), + "ShakaPackager": ( + "shaka-packager", + "packager", + f"packager-{__shaka_platform}", + f"packager-{__shaka_platform}-arm64", + f"packager-{__shaka_platform}-x64", + ), + "CCExtractor": ("ccextractor", "ccextractorwin", "ccextractorwinfull"), + "HolaProxy": ("hola-proxy",), + "MPV": ("mpv",), + "Caddy": ("caddy",), + "MKVToolNix": ("mkvmerge",), + "Mkvpropedit": ("mkvpropedit",), + "DoviTool": ("dovi_tool",), + "HDR10PlusTool": ("hdr10plus_tool", "HDR10Plus_tool"), + "Mp4decrypt": ("mp4decrypt",), + "Docker": ("docker",), + "ML_Worker": ("ML-Worker",), + "Git": ("git",), +} + + +def __getattr__(name: str) -> Optional[Path]: + candidates = __binaries.get(name) + if candidates is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + # Cache the result (including None for an absent binary) so later access is a plain attr hit. + resolved = find(*candidates) + globals()[name] = resolved + return resolved __all__ = ( diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index cc2c84aa..ab49235b 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -586,6 +586,9 @@ def requests( total_bytes = 0 start_time = time.time() last_speed_report = start_time + last_hedge_check = 0.0 + hedge_median = 0.0 # cached; recomputed only when seg_durations grows + hedge_median_len = 0 pool = ThreadPoolExecutor(max_workers=max_workers) event_queue: Queue[dict[str, Any]] = Queue() @@ -672,9 +675,20 @@ def requests( yield dict(downloaded=f"{filesize.decimal(download_speed)}/s") last_speed_report = now - # hedge stuck segments once spare workers exist (pending < max_workers) - if len(pending) < max_workers and seg_durations and not DOWNLOAD_CANCELLED.is_set(): - threshold = max(HEDGE_FACTOR * statistics.median(seg_durations), HEDGE_MIN_WAIT) + # hedge stuck segments once spare workers exist (pending < max_workers); + # throttled to ~0.5s and median recomputed only when seg_durations grows + if ( + len(pending) < max_workers + and seg_durations + and now - last_hedge_check > 0.5 + and not DOWNLOAD_CANCELLED.is_set() + ): + last_hedge_check = now + cur_len = len(seg_durations) # atomic read; appended under seg_lock by workers + if cur_len != hedge_median_len: + hedge_median = statistics.median(seg_durations) + hedge_median_len = cur_len + threshold = max(HEDGE_FACTOR * hedge_median, HEDGE_MIN_WAIT) with seg_lock: stuck = [ i diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index 3aeee536..5f773760 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -8,7 +8,7 @@ import re import shutil import sys from copy import deepcopy -from functools import partial +from functools import lru_cache, partial from typing import Any, Callable, Optional, Union from urllib.parse import urljoin, urlparse from uuid import UUID @@ -1103,11 +1103,19 @@ class DASH: m = re.findall(r"([\d.]+.)", d) return sum(float(x[0:-1]) * {"H": 60 * 60, "M": 60, "S": 1}[x[-1].upper()] for x in m) + @staticmethod + @lru_cache(maxsize=None) + def _field_format_pattern(field: str) -> re.Pattern: + # printf-style `$Field%fmt$` matcher, compiled once per field name + return re.compile(rf"\${re.escape(field)}%([a-z0-9]+)\$", flags=re.I) + @staticmethod def replace_fields(url: str, **kwargs: Any) -> str: + if "$" not in url: + return url for field, value in kwargs.items(): url = url.replace(f"${field}$", str(value)) - m = re.search(rf"\${re.escape(field)}%([a-z0-9]+)\$", url, flags=re.I) + m = DASH._field_format_pattern(field).search(url) if m: url = url.replace(m.group(), f"{value:{m.group(1)}}") return url diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index 74428f07..0c0a41d6 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -1106,7 +1106,7 @@ class HLS: with open(save_path, "wb") as output_file: for segment in segments: with open(segment, "rb") as segment_file: - output_file.write(segment_file.read()) + shutil.copyfileobj(segment_file, output_file, 1024 * 1024) cleanup_segments_and_dirs() return save_path.stat().st_size diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index abfb46fd..30613a3f 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -723,7 +723,16 @@ class Track: return init_data - def repackage(self) -> None: + def repackage(self, bsf_v: Optional[str] = None) -> bool: + """Remux the track with ffmpeg ``-c copy``. + + When ``bsf_v`` is given it is folded into the same pass as ``-bsf:v`` (used to + normalize video VUI colour metadata without a second full-file remux). Repackaging + is mandatory; the bitstream filter is best-effort. If the combined run fails (and it + isn't the AAC-retry case) it is retried once without ``bsf_v`` so the remux still + succeeds. Returns True if the requested ``bsf_v`` was applied (always False when + ``bsf_v`` is None, since nothing was requested). + """ if not self.path or not self.path.exists(): raise ValueError("Cannot repackage a Track that has not been downloaded.") @@ -733,7 +742,7 @@ class Track: original_path = self.path output_path = original_path.with_stem(f"{original_path.stem}_repack") - def _ffmpeg(extra_args: list[str] = None): + def _ffmpeg(extra_args: list[str] = None, bsf: Optional[str] = None): args = [ binaries.FFMPEG, "-hide_banner", @@ -767,9 +776,11 @@ class Track: "bitexact", # only have minimal tag data, reproducible mux "-codec", "copy", - str(output_path), ] ) + if bsf: + args.extend(["-bsf:v", bsf]) + args.append(str(output_path)) subprocess.run( args, @@ -785,12 +796,20 @@ class Track: context={"track_type": self.__class__.__name__, "id": self.id, "file": original_path.name}, ) + bsf_applied = False try: - _ffmpeg() + _ffmpeg(bsf=bsf_v) + bsf_applied = bsf_v is not None except subprocess.CalledProcessError as e: if b"Malformed AAC bitstream detected" in e.stderr: # e.g., TruTV's dodgy encodes - _ffmpeg(["-y", "-bsf:a", "aac_adtstoasc"]) + _ffmpeg(["-y", "-bsf:a", "aac_adtstoasc"], bsf=bsf_v) + bsf_applied = bsf_v is not None + elif bsf_v is not None: + # Repack is mandatory, the VUI bitstream filter is best-effort: retry + # without it so the remux still succeeds; caller falls back to normalize_vui. + output_path.unlink(missing_ok=True) + _ffmpeg() else: raise @@ -808,6 +827,7 @@ class Track: "output_size": output_path.stat().st_size if output_path.exists() else 0, }, ) + return bsf_applied __all__ = ("Track", "DownloadContext") diff --git a/unshackle/core/tracks/video.py b/unshackle/core/tracks/video.py index a3804a41..ead971fc 100644 --- a/unshackle/core/tracks/video.py +++ b/unshackle/core/tracks/video.py @@ -410,20 +410,19 @@ class Video(Track): self.path = output_path original_path.unlink() - def normalize_vui(self) -> bool: - """Rewrite SPS VUI colour metadata to match ``self.range``. + def vui_bsf(self) -> Optional[str]: + """Return the ``-bsf:v`` value that rewrites SPS VUI colour metadata to match ``self.range``. - Some services ship HDR10/HLG bitstreams with stale BT.709 VUI, which makes - downstream tools mis-classify the file. The manifest-derived range is the - source of truth. Skips SDR, DV, and HYBRID. Returns True if the bitstream - was rewritten. + Returns None when no rewrite is needed: SDR/DV/HYBRID ranges, non-AVC/HEVC codecs, a + missing file, or when the bitstream already ships the correct colour metadata. Otherwise + returns a ``{h264,hevc}_metadata=...`` string suitable for ffmpeg ``-bsf:v``. """ if not self.path or not self.path.exists(): - return False + return None if self.codec not in (Video.Codec.AVC, Video.Codec.HEVC): - return False + return None if self.range in (Video.Range.SDR, Video.Range.DV, Video.Range.HYBRID): - return False + return None vui = { Video.Range.HDR10: (9, 16, 9), @@ -431,7 +430,7 @@ class Video(Track): Video.Range.HLG: (9, 18, 9), }.get(self.range) if not vui: - return False + return None # Skip the rewrite when the bitstream VUI already matches the target — avoids a full # bitstream pass on services that already ship correct colour metadata @@ -462,19 +461,32 @@ class Video(Track): vals = dict(line.split("=", 1) for line in probe.stdout.splitlines() if "=" in line) current = (vals.get("color_primaries"), vals.get("color_transfer"), vals.get("color_space")) if current == expected: - return False - - if not binaries.FFMPEG: - raise EnvironmentError('FFmpeg executable "ffmpeg" was not found but is required for this call.') + return None primaries, transfer, matrix = vui filter_key = {Video.Codec.AVC: "h264_metadata", Video.Codec.HEVC: "hevc_metadata"}[self.codec] - bsf = ( + return ( f"{filter_key}=colour_primaries={primaries}" f":transfer_characteristics={transfer}" f":matrix_coefficients={matrix}" ) + def normalize_vui(self) -> bool: + """Rewrite SPS VUI colour metadata to match ``self.range``. + + Some services ship HDR10/HLG bitstreams with stale BT.709 VUI, which makes + downstream tools mis-classify the file. The manifest-derived range is the + source of truth. Skips SDR, DV, and HYBRID. Returns True if the bitstream + was rewritten. + """ + bsf = self.vui_bsf() + if bsf is None: + return False + assert self.path is not None # vui_bsf() returns None when path is unset + + if not binaries.FFMPEG: + raise EnvironmentError('FFmpeg executable "ffmpeg" was not found but is required for this call.') + original_path = self.path output_path = original_path.with_stem(f"{original_path.stem}_vui") try: @@ -508,7 +520,7 @@ class Video(Track): context={ "codec": str(self.codec), "range": str(self.range), - "vui": {"primaries": primaries, "transfer": transfer, "matrix": matrix}, + "vui": bsf, "id": self.id, }, )