From 8bf5444ae96c505a7d6ed393d52fba1f47a685c2 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Sat, 4 Jul 2026 11:38:19 -0600 Subject: [PATCH] refactor(core)!: bundle download_track args into DownloadContext BREAKING CHANGE: HLS/DASH/ISM.download_track now take (track, ctx); external wrappers/monkeypatches must update to the new signature. --- unshackle/core/manifests/dash.py | 28 +++++--------- unshackle/core/manifests/hls.py | 30 ++++++--------- unshackle/core/manifests/ism.py | 31 ++++++---------- unshackle/core/tracks/__init__.py | 15 +++++++- unshackle/core/tracks/track.py | 61 +++++++++++++++++-------------- 5 files changed, 78 insertions(+), 87 deletions(-) diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index 48dabb0d..1d15cf5b 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -9,7 +9,6 @@ import shutil import sys from copy import deepcopy from functools import partial -from pathlib import Path from typing import Any, Callable, Optional, Union from urllib.parse import urljoin, urlparse from uuid import UUID @@ -28,7 +27,7 @@ from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, from unshackle.core.drm import DRM_T, ClearKeyCENC, PlayReady, Widevine from unshackle.core.events import events from unshackle.core.session import RnetSession -from unshackle.core.tracks import Audio, Subtitle, Tracks, Video +from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video from unshackle.core.utilities import is_close_match, log_event, try_ensure_utf8 from unshackle.core.utils.redact import safe_display_url from unshackle.core.utils.xml import load_xml @@ -275,22 +274,15 @@ class DASH: return tracks @staticmethod - def download_track( - track: AnyTrack, - save_path: Path, - save_dir: Path, - progress: partial, - session: Optional[Session] = None, - proxy: Optional[str] = None, - max_workers: Optional[int] = None, - license_widevine: Optional[Callable] = None, - *, - cdm: Optional[object] = None, - ): - if not session: - session = Session() - elif not isinstance(session, (Session, RnetSession)): - raise TypeError(f"Expected session to be a {Session} or {RnetSession}, not {session!r}") + def download_track(track: AnyTrack, ctx: DownloadContext) -> None: + session = ctx.ensure_session() + save_path = ctx.save_path + save_dir = ctx.save_dir + progress = ctx.progress + proxy = ctx.proxy + max_workers = ctx.max_workers + license_widevine = ctx.license_widevine + cdm = ctx.cdm if proxy: session.proxies.update({"all": proxy}) diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index eacad66f..2b1c2eec 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -10,9 +10,8 @@ import shutil import subprocess import sys import time -from functools import partial from pathlib import Path -from typing import Any, Callable, Optional, Union +from typing import Any, Optional, Union from urllib.parse import urljoin from uuid import UUID from zlib import crc32 @@ -33,7 +32,7 @@ from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, from unshackle.core.drm import DRM_T, ClearKey, MonaLisa, PlayReady, Widevine from unshackle.core.events import events from unshackle.core.session import RnetResponse, RnetSession -from unshackle.core.tracks import Audio, Subtitle, Tracks, Video +from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video from unshackle.core.utilities import get_extension, is_close_match, log_event, try_ensure_utf8 from unshackle.core.utils.redact import safe_display_url from unshackle.core.utils.subprocess import log_tool_run @@ -517,22 +516,15 @@ class HLS: return None @staticmethod - def download_track( - track: AnyTrack, - save_path: Path, - save_dir: Path, - progress: partial, - session: Optional[Union[Session, RnetSession]] = None, - proxy: Optional[str] = None, - max_workers: Optional[int] = None, - license_widevine: Optional[Callable] = None, - *, - cdm: Optional[object] = None, - ) -> None: - if not session: - session = Session() - elif not isinstance(session, (Session, RnetSession)): - raise TypeError(f"Expected session to be a {Session} or {RnetSession}, not {session!r}") + def download_track(track: AnyTrack, ctx: DownloadContext) -> None: + session = ctx.ensure_session() + save_path = ctx.save_path + save_dir = ctx.save_dir + progress = ctx.progress + proxy = ctx.proxy + max_workers = ctx.max_workers + license_widevine = ctx.license_widevine + cdm = ctx.cdm if proxy: # Handle proxies differently based on session type diff --git a/unshackle/core/manifests/ism.py b/unshackle/core/manifests/ism.py index a7ffa8dd..0c8913ef 100644 --- a/unshackle/core/manifests/ism.py +++ b/unshackle/core/manifests/ism.py @@ -6,9 +6,7 @@ import html import shutil import struct import urllib.parse -from functools import partial -from pathlib import Path -from typing import Any, Callable, Optional, Union +from typing import Any, Optional, Union import requests from langcodes import Language, tag_is_valid @@ -23,7 +21,7 @@ from unshackle.core.events import events from unshackle.core.manifests.ism_init import (build_init_segment, parse_codec_private_data_colour, read_per_sample_iv_size, read_track_id) from unshackle.core.session import RnetSession -from unshackle.core.tracks import Audio, Subtitle, Track, Tracks, Video +from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Track, Tracks, Video from unshackle.core.utilities import log_event, try_ensure_utf8 from unshackle.core.utils.redact import safe_display_url from unshackle.core.utils.xml import load_xml @@ -365,22 +363,15 @@ class ISM: return tracks @staticmethod - def download_track( - track: AnyTrack, - save_path: Path, - save_dir: Path, - progress: partial, - session: Optional[Union[Session, RnetSession]] = None, - proxy: Optional[str] = None, - max_workers: Optional[int] = None, - license_widevine: Optional[Callable] = None, - *, - cdm: Optional[object] = None, - ) -> None: - if not session: - session = Session() - elif not isinstance(session, (Session, RnetSession)): - raise TypeError(f"Expected session to be a {Session} or {RnetSession}, not {session!r}") + def download_track(track: AnyTrack, ctx: DownloadContext) -> None: + session = ctx.ensure_session() + save_path = ctx.save_path + save_dir = ctx.save_dir + progress = ctx.progress + proxy = ctx.proxy + max_workers = ctx.max_workers + license_widevine = ctx.license_widevine + cdm = ctx.cdm if proxy: session.proxies.update({"all": proxy}) diff --git a/unshackle/core/tracks/__init__.py b/unshackle/core/tracks/__init__.py index 37aca7fa..4d9ed09f 100644 --- a/unshackle/core/tracks/__init__.py +++ b/unshackle/core/tracks/__init__.py @@ -4,8 +4,19 @@ from .chapter import Chapter from .chapters import Chapters from .hybrid import Hybrid from .subtitle import Subtitle -from .track import Track +from .track import DownloadContext, Track from .tracks import Tracks from .video import Video -__all__ = ("Audio", "Attachment", "Chapter", "Chapters", "Hybrid", "Subtitle", "Track", "Tracks", "Video") +__all__ = ( + "Audio", + "Attachment", + "Chapter", + "Chapters", + "DownloadContext", + "Hybrid", + "Subtitle", + "Track", + "Tracks", + "Video", +) diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index 2b364187..11e09c78 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -6,6 +6,7 @@ import shutil import subprocess from collections import defaultdict from copy import copy +from dataclasses import dataclass from enum import Enum from functools import partial from pathlib import Path @@ -47,6 +48,29 @@ def direct_session(session: Union[Session, "RnetSession"]) -> Session: return new +@dataclass +class DownloadContext: + """Shared arguments passed to each manifest's ``download_track``.""" + + save_path: Path + save_dir: Path + progress: partial + session: Optional[Union[Session, "RnetSession"]] = None + proxy: Optional[str] = None + max_workers: Optional[int] = None + license_widevine: Optional[Callable] = None + cdm: Optional[object] = None + + def ensure_session(self) -> Union[Session, "RnetSession"]: + """Return the session, or a new ``Session`` if none was set.""" + session = self.session + if not session: + session = Session() + elif not isinstance(session, (Session, RnetSession)): + raise TypeError(f"Expected session to be a {Session} or {RnetSession}, not {session!r}") + return session + + class Track: class Descriptor(Enum): URL = 1 # Direct URL, nothing fancy @@ -256,33 +280,13 @@ class Track: cleanup() try: - if self.descriptor == self.Descriptor.HLS: - HLS.download_track( - track=self, - save_path=save_path, - save_dir=save_dir, - progress=progress, - session=dl_session, - proxy=proxy, - max_workers=max_workers, - license_widevine=prepare_drm, - cdm=cdm, - ) - elif self.descriptor == self.Descriptor.DASH: - DASH.download_track( - track=self, - save_path=save_path, - save_dir=save_dir, - progress=progress, - session=dl_session, - proxy=proxy, - max_workers=max_workers, - license_widevine=prepare_drm, - cdm=cdm, - ) - elif self.descriptor == self.Descriptor.ISM: - ISM.download_track( - track=self, + manifest_parsers = { + self.Descriptor.HLS: HLS, + self.Descriptor.DASH: DASH, + self.Descriptor.ISM: ISM, + } + if self.descriptor in manifest_parsers: + ctx = DownloadContext( save_path=save_path, save_dir=save_dir, progress=progress, @@ -292,6 +296,7 @@ class Track: license_widevine=prepare_drm, cdm=cdm, ) + manifest_parsers[self.descriptor].download_track(track=self, ctx=ctx) elif self.descriptor == self.Descriptor.URL: try: if not self.drm and track_type in ("Video", "Audio"): @@ -805,4 +810,4 @@ class Track: ) -__all__ = ("Track",) +__all__ = ("Track", "DownloadContext")