From 7898ff6dd605e413fba73bfe78447245009bd895 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Fri, 30 Oct 2026 20:51:21 -0600 Subject: [PATCH] refactor(types): bring mypy to zero errors across unshackle Fix the bugs the type errors exposed (RnetSession.send body, HLS proxy session, title cache fallback, mux progress) and gate mypy in pre-commit, warn-only in CI. --- .github/workflows/ci.yml | 5 + .pre-commit-config.yaml | 9 + pyproject.toml | 2 + tests/core/test_import_track_keys.py | 18 +- unshackle/commands/dl.py | 156 +++++++++++------- unshackle/commands/env.py | 7 +- unshackle/commands/kv.py | 1 + unshackle/commands/search.py | 2 +- unshackle/commands/serve.py | 6 +- unshackle/commands/setup.py | 2 +- unshackle/commands/util.py | 5 +- unshackle/commands/wvd.py | 6 +- unshackle/core/__main__.py | 2 +- unshackle/core/api/compression.py | 3 +- unshackle/core/api/download_manager.py | 30 ++-- unshackle/core/api/errors.py | 2 +- unshackle/core/api/handlers.py | 40 +++-- unshackle/core/binaries.py | 6 +- unshackle/core/cacher.py | 3 +- unshackle/core/cdm/custom_remote_cdm.py | 36 ++-- unshackle/core/cdm/decrypt_labs_remote_cdm.py | 27 +-- unshackle/core/cdm/detect.py | 6 +- unshackle/core/cdm/monalisa/monalisa_cdm.py | 97 ++++++----- unshackle/core/commands.py | 6 +- unshackle/core/config.py | 6 +- unshackle/core/console.py | 20 +-- unshackle/core/constants.py | 9 +- unshackle/core/credential.py | 2 +- unshackle/core/downloaders/requests.py | 30 ++-- unshackle/core/drm/blowfish_stripe.py | 2 +- unshackle/core/drm/clearkey.py | 6 +- unshackle/core/drm/clearkey_cenc.py | 5 +- unshackle/core/drm/monalisa.py | 2 +- unshackle/core/drm/playready.py | 11 +- unshackle/core/drm/playready_errors.py | 2 +- unshackle/core/drm/widevine.py | 11 +- unshackle/core/events.py | 4 +- unshackle/core/import_service.py | 13 +- unshackle/core/manifests/dash.py | 32 ++-- unshackle/core/manifests/hls.py | 51 +++--- unshackle/core/manifests/ism.py | 3 +- unshackle/core/manifests/ism_init.py | 5 +- unshackle/core/music/renderer.py | 6 +- unshackle/core/music/tagger.py | 11 +- unshackle/core/proxies/basic.py | 2 +- unshackle/core/proxies/controld.py | 6 +- unshackle/core/proxies/expressvpn.py | 2 +- unshackle/core/proxies/gluetun.py | 24 +-- unshackle/core/proxies/hola.py | 6 +- unshackle/core/proxies/nordvpn.py | 3 +- unshackle/core/proxies/proton.py | 5 +- unshackle/core/proxies/proxy.py | 4 +- unshackle/core/proxies/surfsharkvpn.py | 2 +- unshackle/core/proxies/windscribevpn.py | 2 +- unshackle/core/remote_service.py | 19 ++- unshackle/core/search_result.py | 2 +- unshackle/core/service.py | 18 +- unshackle/core/service_repo.py | 2 +- unshackle/core/services.py | 14 +- unshackle/core/session.py | 32 ++-- unshackle/core/title_cacher.py | 15 +- unshackle/core/titles/episode.py | 4 +- unshackle/core/titles/movie.py | 4 +- unshackle/core/titles/music.py | 4 +- unshackle/core/titles/title.py | 6 +- unshackle/core/tracks/attachment.py | 10 +- unshackle/core/tracks/audio.py | 6 +- unshackle/core/tracks/chapter.py | 11 +- unshackle/core/tracks/chapters.py | 2 +- unshackle/core/tracks/dv_fixup.py | 1 + unshackle/core/tracks/hybrid.py | 26 +-- unshackle/core/tracks/subtitle.py | 30 ++-- unshackle/core/tracks/track.py | 70 ++++---- unshackle/core/tracks/tracks.py | 47 +++--- unshackle/core/tracks/video.py | 15 +- unshackle/core/utilities.py | 61 ++++--- unshackle/core/utils/bitrate.py | 11 +- unshackle/core/utils/click_types.py | 15 +- unshackle/core/utils/dovi.py | 4 +- unshackle/core/utils/gen_esn.py | 4 +- unshackle/core/utils/redact.py | 6 +- unshackle/core/utils/selector.py | 64 +++++-- unshackle/core/utils/sslciphers.py | 8 +- unshackle/core/utils/subprocess.py | 2 +- unshackle/core/utils/template_formatter.py | 2 +- unshackle/core/vault.py | 4 +- unshackle/core/vaults.py | 4 +- unshackle/services/EXAMPLE/__init__.py | 2 +- unshackle/services/MUSIC_EXAMPLE/__init__.py | 2 +- unshackle/utils/base62.py | 14 +- unshackle/vaults/API.py | 2 +- unshackle/vaults/HTTP.py | 50 +++--- unshackle/vaults/MySQL.py | 23 ++- unshackle/vaults/SQLite.py | 8 +- uv.lock | 11 ++ 95 files changed, 810 insertions(+), 589 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 80fdcdf8..73936bcc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -35,6 +35,11 @@ jobs: run: uv run ruff format --check . - name: Import order run: uv run isort --check-only unshackle tests --skip-glob "*/services/*" + - name: Type check + run: | + out=$(uv run mypy unshackle) || true + echo "$out" + echo "$out" | sed -nE 's/^([^:]+):([0-9]+): error: (.*)$/::warning file=\1,line=\2::\3/p' test: strategy: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d5cb2835..be522d47 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -33,3 +33,12 @@ repos: - id: end-of-file-fixer - id: trailing-whitespace args: [--markdown-linebreak-ext=md] + - repo: local + hooks: + - id: mypy + name: mypy + entry: uv run mypy unshackle + language: system + types: [python] + files: ^unshackle/ + pass_filenames: false diff --git a/pyproject.toml b/pyproject.toml index c967b56f..3de43c31 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,6 +94,7 @@ dev = [ "types-requests>=2.31.0.20240406,<3", "isort>=5.13.2,<8", "ruff>=0.15.12,<0.16", + "types-pyyaml>=6.0,<7", ] test = [ "pytest>=8.3.0,<9", @@ -160,6 +161,7 @@ line_length = 118 check_untyped_defs = true disallow_incomplete_defs = true explicit_package_bases = true +exclude = ["^unshackle/services/_repos/"] disallow_untyped_defs = true follow_imports = "silent" ignore_missing_imports = true diff --git a/tests/core/test_import_track_keys.py b/tests/core/test_import_track_keys.py index f8c373c9..d2ab7b10 100644 --- a/tests/core/test_import_track_keys.py +++ b/tests/core/test_import_track_keys.py @@ -19,7 +19,7 @@ import click from pywidevine.pssh import PSSH from unshackle.commands.dl import dl -from unshackle.core.drm import Widevine +from unshackle.core.drm import ClearKey, Widevine from unshackle.core.drm.key_args import ZERO_KID from unshackle.core.import_service import ImportService from unshackle.core.manifests import DASH, HLS @@ -260,3 +260,19 @@ def test_a_subtitle_that_lists_its_kid_keeps_its_key(tmp_path: Path) -> None: drm = hls_import(tmp_path, subtitle_drm=widevine(SUB_KID, SUB_KEY)) assert drm["es"] and drm["es"][0].content_keys == {SUB_KID: SUB_KEY} + + +def test_a_track_whose_drm_holds_its_own_key_is_left_alone() -> None: + """HLS AES-128 has no content_keys to fill, so injecting the pool into it must not crash.""" + svc = ImportService.__new__(ImportService) + track = Video(url="https://example.invalid/v.m3u8", language="en", drm=[ClearKey(key=VIDEO_KEY)]) + title = SimpleNamespace(id="movie-1", tracks=[track]) + with ( + patch.object(ImportService, "key_pool", return_value=POOL), + patch.object(ImportService, "exported_drm_system", return_value="widevine"), + patch.object(ImportService, "exported_kids", return_value=[VIDEO_KID.hex]), + patch.object(ImportService, "exported_keys", return_value={VIDEO_KID: VIDEO_KEY}), + ): + svc.resolve_server_keys(cast(Any, title)) + + assert isinstance(track.drm[0], ClearKey) diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index e87fa8e3..6578cc34 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -22,7 +22,7 @@ from io import StringIO from itertools import product from pathlib import Path from threading import Lock -from typing import Any, Callable, Collection, Optional, Sequence, TypedDict, Union +from typing import Any, Callable, Collection, Optional, Sequence, TypedDict, Union, cast from uuid import UUID import click @@ -60,7 +60,18 @@ from unshackle.core.downloaders import ( retries_exhausted, set_speed_limit, ) -from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine, cbcs, own_kids, real_kids, verify +from unshackle.core.drm import ( + DRM_T, + ClearKey, + ClearKeyCENC, + MonaLisa, + PlayReady, + Widevine, + cbcs, + own_kids, + real_kids, + verify, +) from unshackle.core.events import events from unshackle.core.export_name import ExportNamer, move_export, plural, season_episodes, title_label from unshackle.core.providers.anilist import parse_anilist_ref @@ -501,7 +512,7 @@ def download_tracks_in_passes( DOWNLOAD_CANCELLED.clear() indexed = list(enumerate(tracks)) if skip_subtitle_errors: - primary = [(i, t) for i, t in indexed if not isinstance(t, Subtitle)] + primary: list[tuple[int, AnyTrack]] = [(i, t) for i, t in indexed if not isinstance(t, Subtitle)] skippable = [(i, t) for i, t in indexed if isinstance(t, Subtitle)] else: primary, skippable = indexed, [] @@ -1254,7 +1265,7 @@ class dl: proxy_providers: Optional[list[Any]] = None, *_: Any, **__: Any, - ): + ) -> None: if not ctx.invoked_subcommand: raise ValueError("A subcommand to invoke was not specified, the main code cannot continue.") @@ -1357,7 +1368,7 @@ class dl: ) # Log binary versions for diagnostics - binary_versions = {} + binary_versions: dict[str, Optional[dict[str, Optional[str]]]] = {} for name, binary in [ ("shaka_packager", binaries.ShakaPackager), ("mp4decrypt", binaries.Mp4decrypt), @@ -1668,21 +1679,21 @@ class dl: quality: list[int], vcodec: list[Video.Codec], acodec: list[Audio.Codec], - vbitrate: int, - abitrate: int, + vbitrate: Optional[int], + abitrate: Optional[int], vbitrate_range: Optional[str], abitrate_range: Optional[str], range_: list[Video.Range], - channels: float, + channels: Optional[float], no_atmos: bool, select_titles: bool, wanted: list[str], latest_episode: bool, latest_episodes: Optional[int], lang: list[str], - v_lang: list[str], - a_lang: list[str], - s_lang: list[str], + v_lang: Optional[list[str]], + a_lang: Optional[list[str]], + s_lang: Optional[list[str]], require_audio: list[str], require_video: list[str], require_subs: list[str], @@ -1707,7 +1718,7 @@ class dl: list_: bool, list_titles: bool, skip_dl: bool, - export: bool, + export: Optional[bool], cdm_only: Optional[bool], no_proxy: bool, no_proxy_download: bool, @@ -1761,7 +1772,7 @@ class dl: # a remote session sends --cdm to the server, which can hand it the tracks above its cap if self.cdm_override and isinstance(service, ImportService): self.log.warning("--cdm is ignored: an import uses the keys in the export") - self._remote_service = service if hasattr(service, "_server_cdm") else None + self._remote_service: Any = service if hasattr(service, "_server_cdm") else None start_time = time.time() lang, lang_excl = partition_exclusions(lang) @@ -2110,7 +2121,7 @@ class dl: console.print(Padding(Rule("[rule.text]Select Titles"), (1, 2))) selection_titles = [] - dependencies = {} + dependencies: dict[int, list[int]] = {} original_indices = {} current_season = None @@ -2232,8 +2243,10 @@ class dl: selected.append(candidates[0]) # Built via per-tag by_language calls, so overlapping tags can append # the same track across iterations; dedupe by id. - seen_ids: set[str] = set() - return [t for t in selected if not (t.id in seen_ids or seen_ids.add(t.id))] + unique: dict[str, Audio] = {} + for t in selected: + unique.setdefault(t.id, t) + return list(unique.values()) base_selection = (v_lang, a_lang, s_lang, range_) @@ -2362,8 +2375,8 @@ class dl: delay = random.randint(slow[0], slow[1]) spinner = Spinner("dots", text=f"Delaying by {delay} seconds...") 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...") + for seconds_left in range(delay, 0, -1): + spinner.update(text=f"Delaying by {seconds_left} seconds...") time.sleep(1) with console.status("Subscribing to events...", spinner="dots"): @@ -2579,6 +2592,10 @@ class dl: fsl_excl_r = resolve_sort_langs(fsl_excl, title.language, exact=exact_lang) if isinstance(title, (Movie, Episode)): + resolution: Optional[int] + color_range: Optional[Video.Range] + codec: Optional[Video.Codec] + vlang: Optional[str] if keep_videos and vcodec: title.tracks.select_video(lambda x: x.codec in vcodec) missing_codecs = [c for c in vcodec if not any(x.codec == c for x in title.tracks.videos)] @@ -2609,20 +2626,20 @@ class dl: # In HYBRID mode, only apply bitrate filter to non-DV tracks # DV tracks are kept regardless since they're only used for RPU metadata title.tracks.select_video( - lambda x: x.range == Video.Range.DV or (x.bitrate and x.bitrate // 1000 == vbitrate) + lambda x: x.range == Video.Range.DV or bool(x.bitrate and x.bitrate // 1000 == vbitrate) ) if not any(x.range != Video.Range.DV for x in title.tracks.videos): self.log.error(f"There's no {vbitrate}kbps Video Track...") sys.exit(1) else: - title.tracks.select_video(lambda x: x.bitrate and x.bitrate // 1000 == vbitrate) + title.tracks.select_video(lambda x: bool(x.bitrate and x.bitrate // 1000 == vbitrate)) if not title.tracks.videos: self.log.error(f"There's no {vbitrate}kbps Video Track...") sys.exit(1) if keep_videos and vbitrate_min is not None and vbitrate_max is not None: title.tracks.select_video( - lambda x: x.bitrate and vbitrate_min <= x.bitrate // 1000 <= vbitrate_max + lambda x: bool(x.bitrate and vbitrate_min <= x.bitrate // 1000 <= vbitrate_max) ) if not title.tracks.videos: self.log.error(f"No Video Track in {vbitrate_min}-{vbitrate_max}kbps range...") @@ -2686,7 +2703,7 @@ class dl: sys.exit(1) has_hybrid = any(r == Video.Range.HYBRID for r in range_) - non_hybrid_ranges = [r for r in range_ if r != Video.Range.HYBRID] + non_hybrid_ranges: list[Video.Range] = [r for r in range_ if r != Video.Range.HYBRID] effective_quality = quality if keep_videos and quality: missing_resolutions = [] @@ -2768,7 +2785,7 @@ class dl: hybrid_quality = effective_quality if not hybrid_quality: - best_resolution = max((v.height for v in hybrid_candidate_tracks), default=None) + best_resolution = max((v.height for v in hybrid_candidate_tracks if v.height), default=None) hybrid_quality = [best_resolution] if best_resolution else [] hybrid_selected = [] if hybrid_quality: @@ -2780,14 +2797,16 @@ class dl: if not any(at_resolution(v, res) for v in hybrid_selected): self.log.warning(f"Skipping {res}p: nothing at that resolution can make a hybrid") - non_hybrid_selected: list[Video] = [] + non_hybrid_selected = [] if non_hybrid_ranges and non_hybrid_tracks: if video_multi_lang: - non_hybrid_langs = list(dict.fromkeys(str(v.language) for v in non_hybrid_tracks)) + non_hybrid_langs: list[Optional[str]] = list( + dict.fromkeys(str(v.language) for v in non_hybrid_tracks) + ) else: non_hybrid_langs = [None] for resolution, color_range, codec, vlang in product( - effective_quality or [None], non_hybrid_ranges, vcodec or [None], non_hybrid_langs + effective_quality or (None,), non_hybrid_ranges, vcodec or (None,), non_hybrid_langs ): candidates = [ t @@ -2814,11 +2833,13 @@ class dl: elif keep_videos: selected_videos: list[Video] = [] if video_multi_lang: - unique_video_langs = list(dict.fromkeys(str(v.language) for v in title.tracks.videos)) + unique_video_langs: list[Optional[str]] = list( + dict.fromkeys(str(v.language) for v in title.tracks.videos) + ) else: unique_video_langs = [None] for resolution, color_range, codec, vlang in product( - effective_quality or [None], range_ or [None], vcodec or [None], unique_video_langs + effective_quality or (None,), range_ or (None,), vcodec or (None,), unique_video_langs ): candidates = [ t @@ -2845,7 +2866,7 @@ class dl: self.log.info("No Dolby Vision track: Dolby Vision will be generated from HDR10+") if gate_failure: - other_ranges = [r for r in range_ if r != Video.Range.HYBRID] + other_ranges: list[Video.Range] = [r for r in range_ if r != Video.Range.HYBRID] if best_available and other_ranges: self.log.warning( f"{gate_failure}; continuing with {', '.join(r.name for r in other_ranges)}" @@ -2853,12 +2874,14 @@ class dl: range_ = other_ranges fallback_pool = pre_hybrid_videos if video_multi_lang: - fallback_langs = list(dict.fromkeys(str(v.language) for v in fallback_pool)) + fallback_langs: list[Optional[str]] = list( + dict.fromkeys(str(v.language) for v in fallback_pool) + ) else: fallback_langs = [None] fallback_selected: list[Video] = [] for resolution, color_range, codec, vlang in product( - effective_quality or [None], other_ranges, vcodec or [None], fallback_langs + effective_quality or (None,), other_ranges, vcodec or (None,), fallback_langs ): candidates = [ t @@ -3111,7 +3134,7 @@ class dl: download_table = Table.grid() download_table.add_row(selected_tracks) - def prepare_drm_for(track: AnyTrack) -> Callable: + def prepare_drm_for(track: AnyTrack) -> partial: prepare = partial( partial(self.prepare_drm, table=download_table), track=track, @@ -3343,6 +3366,7 @@ class dl: if sidecar_format == "original": config.directories.temp.mkdir(parents=True, exist_ok=True) for subtitle in sidecar_subtitles: + assert subtitle.path is not None original_path = ( config.directories.temp / f"sidecar_original_{subtitle.id}{subtitle.path.suffix}" ) @@ -3626,6 +3650,7 @@ class dl: if len(range_) > 1: task_description += f" {head.range.name}" if len(vcodec) > 1: + assert head.codec is not None task_description += f" {head.codec.name}" task_tracks = Tracks(title.tracks) + title.tracks.chapters + title.tracks.attachments @@ -3788,7 +3813,7 @@ class dl: sidecar_files[first_path] = list(created) self.log.info(f"Saved {len(created)} sidecar subtitle files") - for muxed_path, base_filename, sidecar_dir in sidecar_targets[1:]: + for sidecar_path, base_filename, sidecar_dir in sidecar_targets[1:]: sidecar_dir.mkdir(parents=True, exist_ok=True) copies = [] for source in created: @@ -3797,7 +3822,7 @@ class dl: shutil.copy2(source, target) copies.append(target) if copies: - sidecar_files[muxed_path] = copies + sidecar_files[sidecar_path] = copies self.log.info(f"Saved {len(copies)} sidecar subtitle files") for track in title.tracks: @@ -3865,11 +3890,11 @@ class dl: track = next((t for t in title.tracks if t.path == track_path), None) if track: if isinstance(track, Video): - track_suffix = f".{track.codec.name if hasattr(track.codec, 'name') else 'video'}" + track_suffix = f".{track.codec.name if track.codec is not None else 'video'}" elif isinstance(track, Audio): lang_suffix = f".{track.language}" if track.language else "" track_suffix = ( - f"{lang_suffix}.{track.codec.name if hasattr(track.codec, 'name') else 'audio'}" + f"{lang_suffix}.{track.codec.name if track.codec is not None else 'audio'}" ) elif isinstance(track, Subtitle): lang_suffix = f".{track.language}" if track.language else "" @@ -4002,6 +4027,7 @@ class dl: The sample starts 2 minutes in. A file shorter than 3 minutes gives its middle 60 seconds. """ duration = float(media_info.general_tracks[0].duration or 0) / 1000 + start: float start, end = DEFAULT_SAMPLE length = end - start if duration < end: @@ -4714,7 +4740,7 @@ class dl: if not content_keys or not self.vaults: return self.vaults.clear_licensed_flags(content_keys) - successful_caches = self.vaults.add_keys(content_keys) + successful_caches = self.vaults.add_keys(cast(dict[Union[UUID, str], str], content_keys)) kids, caches = self.vault_cache_tally or (set(), successful_caches) self.vault_cache_tally = (kids | content_keys.keys(), min(caches, successful_caches)) @@ -4764,7 +4790,7 @@ class dl: cdm = self.get_cdm(self.service, self.profile, drm=system) return cdm - def licence_one_drm(self, prepare: Callable, drm: DRM_T, track_kid: Optional[UUID]) -> None: + def licence_one_drm(self, prepare: Callable, drm: Union[Widevine, PlayReady], track_kid: Optional[UUID]) -> None: """Send one challenge for the KIDs of this DRM, also when the DRM or a vault has the content keys. This is the ``--cdm-only`` path. The challenge goes out the first time a run sees the KIDs @@ -4805,14 +4831,15 @@ class dl: a run. A DRM system that fails stops the title. With ``tolerate`` it is a warning, if the other DRM system gave the content key of the track. """ - if type(drm) not in (Widevine, PlayReady): + if not isinstance(drm, (Widevine, PlayReady)): prepare(drm, track_kid=track_kid) return other_type = PlayReady if type(drm) is Widevine else Widevine - others = [d for d in track.drm or [] if type(d) is other_type] + others = [d for d in track.drm or [] if isinstance(d, (Widevine, PlayReady)) and type(d) is other_type] if other_type is PlayReady and len(others) > 1: - others[0].absorb(*others[1:]) + playready = [d for d in others if isinstance(d, PlayReady)] + playready[0].absorb(*playready[1:]) if not others and other_type.__name__ not in self.all_drm_warned: self.all_drm_warned.add(other_type.__name__) self.log.warning( @@ -4851,14 +4878,14 @@ class dl: def prepare_drm( self, - drm: DRM_T, + drm: Union[ClearKeyCENC, Widevine, PlayReady, MonaLisa], track: AnyTrack, title: Title_T, certificate: Callable, licence: Callable, clearkey_licence: Optional[Callable] = None, track_kid: Optional[UUID] = None, - table: Table = None, + table: Optional[Table] = None, cdm_only: bool = False, vaults_only: bool = False, export: Optional[Path] = None, @@ -4873,6 +4900,7 @@ class dl: """ if not drm: return + assert table is not None svc_for_cdm = getattr(self, "_remote_service", None) server_cdm = getattr(svc_for_cdm, "_server_cdm", getattr(self, "server_cdm", False)) @@ -4916,6 +4944,7 @@ class dl: server_vault_keys = {} if cdm_only else getattr(svc_for_cdm, "server_vault_keys", {}) for kid, key in drm.content_keys.items(): if server_vault_keys.get(kid) == key: + assert svc_for_cdm is not None self.vaults.sources[kid] = (key, svc_for_cdm.server_vault) new_keys = { kid: key @@ -4997,6 +5026,7 @@ class dl: track_cdm = getattr(svc_for_cdm, "local_cdm", None) local_drm = track.get_drm_for_cdm(track_cdm) if local_drm is not None and local_drm is not drm: + assert not isinstance(local_drm, ClearKey) # a mid-download handover leaves drm on the server's DRM system, and the caller decrypts with it self.prepare_drm( local_drm, @@ -5048,7 +5078,7 @@ class dl: }, ) - pending_vault_writes: list[Callable[[], Any]] = [] + pending_vault_writes = [] with self.drm_lock(drm): pssh_display = self.truncate_pssh_for_display(drm.pssh.dumps(), "Widevine") cek_tree = Tree(Text.assemble(("Widevine", "cyan"), (f"({pssh_display})", "text"), overflow="fold")) @@ -5068,7 +5098,7 @@ class dl: is_track_kid = ["", "*"][kid == track_kid] key = drm.content_keys[kid] label = f"[text2]{kid.hex}:{key}{is_track_kid}" - if not any(f"{kid.hex}:{key}" in x.label for x in cek_tree.children): + if not any(f"{kid.hex}:{key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) continue @@ -5078,7 +5108,7 @@ class dl: if cached_key: drm.content_keys[kid] = cached_key label = f"[text2]{kid.hex}:{cached_key}{is_track_kid} from cache" - if not any(f"{kid.hex}:{cached_key}" in x.label for x in cek_tree.children): + if not any(f"{kid.hex}:{cached_key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) log_event( "license_cache_hit", @@ -5098,13 +5128,13 @@ class dl: if content_key: drm.content_keys[kid] = content_key label = f"[text2]{kid.hex}:{content_key}{is_track_kid} from {vault_used}" - if not any(f"{kid.hex}:{content_key}" in x.label for x in cek_tree.children): + if not any(f"{kid.hex}:{content_key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) self.LICENSE_KEY_CACHE[kid] = content_key if self.debug_logger: self.debug_logger.log_vault_query( - vault_name=vault_used, + vault_name=str(vault_used), operation="get_key_success", service=self.service, context={ @@ -5190,7 +5220,7 @@ class dl: key = f"[red]{key}[/]" is_track_kid_marker = ["", "*"][kid_ == track_kid] label = f"[text2]{kid_.hex}:{key}{is_track_kid_marker}" - if not any(f"{kid_.hex}:{key}" in x.label for x in cek_tree.children): + if not any(f"{kid_.hex}:{key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) drm.content_keys = { @@ -5267,7 +5297,7 @@ class dl: is_track_kid = ["", "*"][kid == track_kid] key = drm.content_keys[kid] label = f"[text2]{kid.hex}:{key}{is_track_kid}" - if not any(f"{kid.hex}:{key}" in x.label for x in cek_tree.children): + if not any(f"{kid.hex}:{key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) continue @@ -5277,7 +5307,7 @@ class dl: if cached_key: drm.content_keys[kid] = cached_key label = f"[text2]{kid.hex}:{cached_key}{is_track_kid} from cache" - if not any(f"{kid.hex}:{cached_key}" in x.label for x in cek_tree.children): + if not any(f"{kid.hex}:{cached_key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) log_event( "license_cache_hit", @@ -5297,13 +5327,13 @@ class dl: if content_key: drm.content_keys[kid] = content_key label = f"[text2]{kid.hex}:{content_key}{is_track_kid} from {vault_used}" - if not any(f"{kid.hex}:{content_key}" in x.label for x in cek_tree.children): + if not any(f"{kid.hex}:{content_key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) self.LICENSE_KEY_CACHE[kid] = content_key if self.debug_logger: self.debug_logger.log_vault_query( - vault_name=vault_used, + vault_name=str(vault_used), operation="get_key_success", service=self.service, context={ @@ -5370,7 +5400,7 @@ class dl: for kid_, key in drm.content_keys.items(): is_track_kid_marker = ["", "*"][kid_ == track_kid] label = f"[text2]{kid_.hex}:{key}{is_track_kid_marker}" - if not any(f"{kid_.hex}:{key}" in x.label for x in cek_tree.children): + if not any(f"{kid_.hex}:{key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) drm.content_keys.update(from_vaults) @@ -5420,7 +5450,7 @@ class dl: is_track_kid = ["", "*"][kid == track_kid] key = drm.content_keys[kid] label = f"[text2]{kid.hex}:{key}{is_track_kid}" - if not any(f"{kid.hex}:{key}" in x.label for x in cek_tree.children): + if not any(f"{kid.hex}:{key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) continue @@ -5430,7 +5460,7 @@ class dl: if cached_key: drm.content_keys[kid] = cached_key label = f"[text2]{kid.hex}:{cached_key}{is_track_kid} from cache" - if not any(f"{kid.hex}:{cached_key}" in x.label for x in cek_tree.children): + if not any(f"{kid.hex}:{cached_key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) log_event( "license_cache_hit", @@ -5450,7 +5480,7 @@ class dl: if content_key: drm.content_keys[kid] = content_key label = f"[text2]{kid.hex}:{content_key}{is_track_kid} from {vault_used}" - if not any(f"{kid.hex}:{content_key}" in x.label for x in cek_tree.children): + if not any(f"{kid.hex}:{content_key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) self.LICENSE_KEY_CACHE[kid] = content_key elif vaults_only: @@ -5511,7 +5541,7 @@ class dl: for kid_, key in drm.content_keys.items(): is_track_kid_marker = ["", "*"][kid_ == track_kid] label = f"[text2]{kid_.hex}:{key}{is_track_kid_marker}" - if not any(f"{kid_.hex}:{key}" in x.label for x in cek_tree.children): + if not any(f"{kid_.hex}:{key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) drm.content_keys.update(from_vaults) @@ -5554,7 +5584,7 @@ class dl: for kid_, key in drm.content_keys.items(): label = f"[text2]{kid_.hex}:{key}*" - if not any(f"{kid_.hex}:{key}" in x.label for x in cek_tree.children): + if not any(f"{kid_.hex}:{key}" in str(x.label) for x in cek_tree.children): cek_tree.add(label) if cek_tree.children and not pre_existing_tree: @@ -5577,6 +5607,7 @@ class dl: return profile_cookie_file elif default_cookie_file.exists(): return default_cookie_file + return None @staticmethod def get_cookie_jar(service: str, profile: Optional[str]) -> Optional[CookieJar]: @@ -5639,7 +5670,7 @@ class dl: return cookie_jar @staticmethod - def save_cookies(path: Path, cookies: CookieJar): + def save_cookies(path: Path, cookies: CookieJar) -> None: if hasattr(cookies, "jar"): cookies = cookies.jar @@ -5664,6 +5695,7 @@ class dl: if isinstance(credentials, list): return Credential(*credentials) return Credential.loads(credentials) # type: ignore + return None def cdm_from_service_entry(self) -> bool: """Whether the service's own ``cdm`` entry named the loaded device, not ``--cdm`` or the ``default`` entry.""" @@ -5679,7 +5711,7 @@ class dl: profile: Optional[str] = None, drm: Optional[str] = None, quality: Optional[int] = None, - ) -> Optional[object]: + ) -> Any: """ Get CDM for a specified service (either Local or Remote CDM). Now supports quality-based selection when the caller gives a quality. @@ -5705,7 +5737,7 @@ class dl: ): quality_keys.append(key) - def sort_quality_key(key): + def sort_quality_key(key: str) -> tuple[int, int]: if key.isdigit(): return (0, int(key)) # Exact matches first elif key.startswith(">="): diff --git a/unshackle/commands/env.py b/unshackle/commands/env.py index 9ed65126..fdaf55de 100644 --- a/unshackle/commands/env.py +++ b/unshackle/commands/env.py @@ -367,7 +367,7 @@ def info() -> None: table.add_column("Path", style="text2", no_wrap=False, overflow="fold") path_vars = { - x: Path(os.getenv(x)) + x: Path(os.environ[x]) for x in ("TEMP", "APPDATA", "LOCALAPPDATA", "USERPROFILE") if sys.platform == "win32" and os.getenv(x) } @@ -386,11 +386,12 @@ def info() -> None: table.add_row(name.title(), paths_str) else: path = attr_value.resolve() + shown = str(path) for var, var_path in path_vars.items(): if path.is_relative_to(var_path): - path = rf"%{var}%\{path.relative_to(var_path)}" + shown = rf"%{var}%\{path.relative_to(var_path)}" break - table.add_row(name.title(), str(path)) + table.add_row(name.title(), shown) print_wide(table) diff --git a/unshackle/commands/kv.py b/unshackle/commands/kv.py index d2906f87..a460b9c3 100644 --- a/unshackle/commands/kv.py +++ b/unshackle/commands/kv.py @@ -256,6 +256,7 @@ def sync( if not len(vaults) > 1: raise click.ClickException("You must provide more than one Vault to sync.") + assert copy.callback is not None ctx.invoke( copy.callback, to_vault_name=vaults[0], diff --git a/unshackle/commands/search.py b/unshackle/commands/search.py index 1eb00c0f..1701f2f4 100644 --- a/unshackle/commands/search.py +++ b/unshackle/commands/search.py @@ -73,7 +73,7 @@ def search( local_proxy: bool = False, remote: bool = False, server: Optional[str] = None, -): +) -> None: if not ctx.invoked_subcommand: raise ValueError("A subcommand to invoke was not specified, the main code cannot continue.") diff --git a/unshackle/commands/serve.py b/unshackle/commands/serve.py index 5c370cf8..0a2aedbc 100644 --- a/unshackle/commands/serve.py +++ b/unshackle/commands/serve.py @@ -297,6 +297,8 @@ def _install_config_reload(app: web.Application, *, no_key: bool, serve_widevine reload = partial( reload_config, _app, no_key=no_key, serve_widevine=serve_widevine, serve_playready=serve_playready ) + if sys.platform == "win32": + return try: asyncio.get_running_loop().add_signal_handler(signal.SIGHUP, reload) except (AttributeError, NotImplementedError): @@ -536,9 +538,9 @@ def serve( cdm_parts.append("pyplayready CDM") log.info(f"Starting integrated server ({' + '.join(cdm_parts)} + REST API)") - def create_serve_authentication(serve_playready_flag: bool): + def create_serve_authentication(serve_playready_flag: bool) -> Any: @web.middleware - async def serve_authentication(request: web.Request, handler) -> web.StreamResponse: + async def serve_authentication(request: web.Request, handler: Any) -> web.StreamResponse: secret_key = request_secret_key(request) if request.path.startswith(DASHBOARD_PREFIX): return await handler(request) diff --git a/unshackle/commands/setup.py b/unshackle/commands/setup.py index c88cefce..65dd195c 100644 --- a/unshackle/commands/setup.py +++ b/unshackle/commands/setup.py @@ -135,7 +135,7 @@ DATA_DIRS = { key: getattr(Config._Directories, key).name for key in ("cookies", "wvds", "prds", "cache", "logs", "exports", "dcsl", "temp") } -DATA_DIRS["services"] = Config._Directories.services[0].name +DATA_DIRS["services"] = Path(Config._Directories.services[0]).name say = partial(console.print, soft_wrap=True) diff --git a/unshackle/commands/util.py b/unshackle/commands/util.py index c5f312a3..f12d040c 100644 --- a/unshackle/commands/util.py +++ b/unshackle/commands/util.py @@ -136,7 +136,7 @@ def crop(path: Path, aspect: str, letter: bool, offset: int, preview: bool) -> N ffmpeg_call = subprocess.Popen( [ - binaries.FFMPEG, + str(binaries.FFMPEG), "-nostdin", "-y", "-i", @@ -227,7 +227,7 @@ def range_(path: Path, full: bool, preview: bool) -> None: ffmpeg_call = subprocess.Popen( [ - binaries.FFMPEG, + str(binaries.FFMPEG), "-nostdin", "-y", "-i", @@ -299,6 +299,7 @@ def test(path: Path, map_: str) -> None: encoding="utf-8", errors="replace", ) + assert p.stderr is not None reached_output = False errors = 0 for line in p.stderr: diff --git a/unshackle/commands/wvd.py b/unshackle/commands/wvd.py index 9a10bc55..57e05ce2 100644 --- a/unshackle/commands/wvd.py +++ b/unshackle/commands/wvd.py @@ -265,8 +265,8 @@ def new( log.info("VMP:") if device.client_id.vmp_data: - file_hashes = FileHashes() - file_hashes.ParseFromString(device.client_id.vmp_data) - log.info(str(file_hashes)) + vmp = FileHashes() + vmp.ParseFromString(device.client_id.vmp_data) + log.info(str(vmp)) else: log.info("None") diff --git a/unshackle/core/__main__.py b/unshackle/core/__main__.py index 89aaaa02..4cd5296a 100644 --- a/unshackle/core/__main__.py +++ b/unshackle/core/__main__.py @@ -112,7 +112,7 @@ def main(version: bool, debug: bool) -> None: @atexit.register -def cleanup(): +def cleanup() -> None: """Clean up resources on exit.""" close_debug_logger() diff --git a/unshackle/core/api/compression.py b/unshackle/core/api/compression.py index 2561d840..bb5a62a9 100644 --- a/unshackle/core/api/compression.py +++ b/unshackle/core/api/compression.py @@ -4,6 +4,7 @@ from __future__ import annotations import gzip import zlib +from typing import Any from aiohttp import web @@ -25,7 +26,7 @@ def safe_inflate(data: bytes, max_size: int = MAX_INFLATE_BYTES) -> bytes: @web.middleware -async def compression_middleware(request: web.Request, handler) -> web.StreamResponse: +async def compression_middleware(request: web.Request, handler: Any) -> web.StreamResponse: """Compress JSON responses with gzip when the client supports it.""" response = await handler(request) diff --git a/unshackle/core/api/download_manager.py b/unshackle/core/api/download_manager.py index 32fff10c..b6f38edf 100644 --- a/unshackle/core/api/download_manager.py +++ b/unshackle/core/api/download_manager.py @@ -15,7 +15,7 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta from enum import Enum from pathlib import Path -from typing import Any, Callable, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, cast from unshackle.core.api.events import bus, publish_service_event from unshackle.core.api.sanitize import sanitize_log @@ -23,6 +23,9 @@ from unshackle.core.utils.click_types import AUDIO_CODEC_LIST, QUALITY_LIST, VID from unshackle.core.utils.redact import REDACTED, URL_USERINFO_RE, redact_all, redact_path, redact_text from unshackle.core.utils.service_params import service_options +if TYPE_CHECKING: + from unshackle.core.service import Service + log = logging.getLogger("download_manager") @@ -145,7 +148,7 @@ class DownloadJob: def to_dict(self, include_full_details: bool = False) -> Dict[str, Any]: """Convert job to dictionary for JSON response.""" - result = { + result: Dict[str, Any] = { "job_id": self.job_id, "status": self.status.value, "created_time": self.created_time.isoformat(), @@ -470,7 +473,7 @@ def perform_download( service_config = load_yaml(service_config_path) else: service_config = {} - merge_dict(config.services.get(service), service_config) + merge_dict(config.services.get(service) or {}, service_config) from unshackle.commands.dl import dl as dl_command @@ -561,7 +564,8 @@ def perform_download( if remote: from unshackle.core.remote_service import make_remote_service - service_instance = make_remote_service(service_ctx, remote, service, title_id, params) + # RemoteService duck-types Service for dl.result() + service_instance = cast("Service", make_remote_service(service_ctx, remote, service, title_id, params)) else: service_module = Services.load(service) @@ -570,7 +574,7 @@ def perform_download( service_init_params = inspect.signature(service_module.__init__).parameters - service_kwargs = {} + service_kwargs: Dict[str, Any] = {} if "title" in service_init_params: service_kwargs["title"] = title_id @@ -704,7 +708,7 @@ def perform_download( class DownloadQueueManager: """Manages download job queue with configurable concurrency limits.""" - def __init__(self, max_concurrent_downloads: int = 2, job_retention_hours: int = 24): + def __init__(self, max_concurrent_downloads: int = 2, job_retention_hours: int = 24) -> None: self.max_concurrent_downloads = max_concurrent_downloads self.job_retention_hours = job_retention_hours @@ -724,7 +728,9 @@ class DownloadQueueManager: f"Initialized download queue manager: max_concurrent={max_concurrent_downloads}, retention_hours={job_retention_hours}" ) - def create_job(self, service: str, title_id: str, owner_key: Optional[str] = None, **parameters) -> DownloadJob: + def create_job( + self, service: str, title_id: str, owner_key: Optional[str] = None, **parameters: Any + ) -> DownloadJob: """Make a new download job and add it to the queue.""" from unshackle.core.api.session_store import refuse_while_draining @@ -911,7 +917,7 @@ class DownloadQueueManager: return len(jobs_to_remove) - async def start_workers(self): + async def start_workers(self) -> None: """Start job worker tasks to process the download queue.""" if self._workers_started: return @@ -927,7 +933,7 @@ class DownloadQueueManager: log.info(f"Started {self.max_concurrent_downloads} download workers") - async def shutdown(self): + async def shutdown(self) -> None: """Shutdown the queue manager and cancel all active downloads.""" log.info("Shutting down download queue manager") self._shutdown_event.set() @@ -971,7 +977,7 @@ class DownloadQueueManager: if self._active_downloads: await asyncio.gather(*self._active_downloads.values(), return_exceptions=True) - async def download_worker(self, worker_name: str): + async def download_worker(self, worker_name: str) -> None: """Job worker task that processes jobs from the queue.""" log.debug(f"Download worker {worker_name} started") @@ -1021,7 +1027,7 @@ class DownloadQueueManager: except Exception as e: log.error(f"Worker {worker_name} error: {e}") - async def execute_download(self, job: DownloadJob): + async def execute_download(self, job: DownloadJob) -> None: """Do the actual download for a job.""" log.info(f"Executing download for job {job.job_id}") @@ -1245,7 +1251,7 @@ class DownloadQueueManager: except OSError: pass - async def cleanup_worker(self): + async def cleanup_worker(self) -> None: """Job worker that periodically cleans up old jobs.""" while not self._shutdown_event.is_set(): try: diff --git a/unshackle/core/api/errors.py b/unshackle/core/api/errors.py index 25f9d962..85f2de22 100644 --- a/unshackle/core/api/errors.py +++ b/unshackle/core/api/errors.py @@ -70,7 +70,7 @@ class APIError(Exception): details: dict[str, Any] | None = None, retryable: bool = False, http_status: int | None = None, - ): + ) -> None: super().__init__(message) self.error_code = error_code self.message = message diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index b80c3984..382bc97a 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -29,7 +29,7 @@ from unshackle.core.api.session_store import SessionStore from unshackle.core.cacher import Cacher from unshackle.core.cdm.detect import cdm_type_stub from unshackle.core.config import config, load_yaml -from unshackle.core.constants import AUDIO_CODEC_MAP, DYNAMIC_RANGE_MAP, VIDEO_CODEC_MAP +from unshackle.core.constants import AUDIO_CODEC_MAP, DYNAMIC_RANGE_MAP, VIDEO_CODEC_MAP, AnyTrack from unshackle.core.drm import BlowfishStripe, cbcs from unshackle.core.providers.anilist import parse_anilist_ref from unshackle.core.proxies.resolve import initialize_proxy_providers, resolve_proxy @@ -199,7 +199,7 @@ def load_service_yaml(normalized_service: str) -> dict: service_config = load_yaml(service_config_path) or {} else: service_config = {} - merge_dict(config.services.get(normalized_service), service_config) + merge_dict(config.services.get(normalized_service) or {}, service_config) return service_config @@ -295,7 +295,7 @@ def server_login_material( cred_data = data.get("credentials") if isinstance(cred_data, dict) and cred_data.get("username") is not None: credential = Credential( - username=cred_data["username"], password=cred_data.get("password"), extra=cred_data.get("extra") + username=cred_data["username"], password=cred_data.get("password") or "", extra=cred_data.get("extra") ) return load_client_cookies(data.get("cookies")), credential @@ -463,7 +463,10 @@ async def relay_target( geofence proxy; ``(None, None, data)`` for a local one.""" if not data.get("remote"): return None, None, data - server, info = await asyncio.to_thread(remote_server, data, request) + relay = await asyncio.to_thread(remote_server, data, request) + if relay is None: + return None, None, data + server, info = relay return server, info, {**data, "service": info["tag"], **geofence_proxy(data, info, request)} @@ -1450,7 +1453,7 @@ def serialize_title(title: Title_T) -> Dict[str, Any]: is_episode = isinstance(title, Episode) is_song = isinstance(title, Song) episode_part = title.part if isinstance(title, Episode) else None - if is_episode: + if isinstance(title, Episode): # no part suffix here: remote_service.build_title rebuilds the Episode from this dict, so a # suffixed name plus the structural `part` below would render the part twice in the filename name = title.name if title.name else f"Episode {title.number:02d}" @@ -1507,7 +1510,7 @@ def stamp_service_flags(serialized: Dict[str, Any], service_instance: Any) -> Di return serialized -def extract_manifests(tracks) -> List[Dict[str, Any]]: +def extract_manifests(tracks: Tracks) -> List[Dict[str, Any]]: """Extract manifest data from tracks for client-side re-parsing. Serializes DASH and ISM manifest XML as zlib-compressed base64 strings @@ -1557,7 +1560,7 @@ def extract_manifests(tracks) -> List[Dict[str, Any]]: return manifests -def extract_track_manifests(tracks) -> List[Dict[str, Any]]: +def extract_track_manifests(tracks: Tracks) -> List[Dict[str, Any]]: """Build a one-AdaptationSet MPD for each track the client cannot find again in the served manifest. A service can build an AdaptationSet of its own, out of merged segment timelines or @@ -1618,7 +1621,7 @@ def extract_track_manifests(tracks) -> List[Dict[str, Any]]: return fragments -def serialize_drm(drm_list) -> Optional[List[Dict[str, Any]]]: +def serialize_drm(drm_list: Any) -> Optional[List[Dict[str, Any]]]: """Serialise DRM objects to JSON-serializable list.""" if not drm_list: return None @@ -1898,8 +1901,8 @@ async def search_handler(data: Dict[str, Any], request: Optional[web.Request] = async def list_titles_handler(data: Dict[str, Any], request: Optional[web.Request] = None) -> web.Response: """Answer the list-titles request.""" require_fields(data, "service", "title_id") - service_tag = data.get("service") - title_id = data.get("title_id") + service_tag = data["service"] + title_id = data["title_id"] profile = client_profile(data) server, info, data = await relay_target(data, request) @@ -1914,7 +1917,7 @@ async def list_titles_handler(data: Dict[str, Any], request: Optional[web.Reques request_cache = RequestCache() lease = None try: - if server: + if server and info: lease, titles = await leased_titles(data, server, info, profile, title_id, request, True) service_instance = lease.service else: @@ -1946,8 +1949,8 @@ async def list_titles_handler(data: Dict[str, Any], request: Optional[web.Reques async def list_tracks_handler(data: Dict[str, Any], request: Optional[web.Request] = None) -> web.Response: """Answer the list-tracks request.""" require_fields(data, "service", "title_id") - service_tag = data.get("service") - title_id = data.get("title_id") + service_tag = data["service"] + title_id = data["title_id"] profile = client_profile(data) server, info, data = await relay_target(data, request) @@ -1962,7 +1965,7 @@ async def list_tracks_handler(data: Dict[str, Any], request: Optional[web.Reques request_cache = RequestCache() lease = None try: - if server: + if server and info: lease, titles = await leased_titles(data, server, info, profile, title_id, request, False) service_instance = lease.service call: Callable[..., Any] = partial(remote_call, lease) @@ -2076,7 +2079,7 @@ async def list_tracks_handler(data: Dict[str, Any], request: Optional[web.Reques continue if episodes_data: - response = {"episodes": episodes_data} + response: Dict[str, Any] = {"episodes": episodes_data} if failed_episodes: response["unavailable_episodes"] = failed_episodes return request_cache.json_response(response) @@ -2430,8 +2433,8 @@ async def download_handler(data: Dict[str, Any], request: Optional[web.Request] from unshackle.core.api.download_manager import get_download_manager require_fields(data, "service", "title_id") - service_tag = data.get("service") - title_id = data.get("title_id") + service_tag = data["service"] + title_id = data["title_id"] server, info, data = await relay_target(data, request) normalized_service = info["tag"] if info else validate_service(service_tag, request) @@ -2544,7 +2547,7 @@ async def list_download_jobs_handler(data: Dict[str, Any], request: Optional[web reverse = sort_order == "desc" - def get_sort_key(job): + def get_sort_key(job: Any) -> Any: """Get the value to sort on, and give a None value a default.""" value = getattr(job, sort_by, None) if value is None: @@ -4170,6 +4173,7 @@ async def session_tracks_handler( session.touch() title_tracks: Dict[str, Any] = {} + track: AnyTrack for track in tracks.videos: title_tracks[str(track.id)] = track session.tracks[str(track.id)] = track diff --git a/unshackle/core/binaries.py b/unshackle/core/binaries.py index 0928a5ff..fc2e6500 100644 --- a/unshackle/core/binaries.py +++ b/unshackle/core/binaries.py @@ -44,9 +44,9 @@ def find(*names: str, search_dirs: Optional[list[Path]] = None) -> Optional[Path return path # Fall back to system PATH - path = shutil.which(name) - if path: - return Path(path) + found = shutil.which(name) + if found: + return Path(found) return None diff --git a/unshackle/core/cacher.py b/unshackle/core/cacher.py index d5b21dd0..86db01e6 100644 --- a/unshackle/core/cacher.py +++ b/unshackle/core/cacher.py @@ -32,7 +32,7 @@ class Cacher: self.service_tag = service_tag self.key = key self.version = version - self.data = data or {} + self.data: Any = data or {} self.expiration = expiration if self.expiration and self.expired: @@ -45,6 +45,7 @@ class Cacher: @property def path(self) -> Path: """Get the path at which the cache will be read and written.""" + assert self.key is not None, "Cacher has no key" return (config.directories.cache / self.service_tag / self.key).with_suffix(".json") @property diff --git a/unshackle/core/cdm/custom_remote_cdm.py b/unshackle/core/cdm/custom_remote_cdm.py index 1a1d4f7a..2946f9cb 100644 --- a/unshackle/core/cdm/custom_remote_cdm.py +++ b/unshackle/core/cdm/custom_remote_cdm.py @@ -17,7 +17,7 @@ from unshackle.core.vaults import Vaults class MockCertificateChain: """Mock certificate chain for PlayReady compatibility.""" - def __init__(self, name: str): + def __init__(self, name: str) -> None: self._name = name def get_name(self) -> str: @@ -33,7 +33,7 @@ class Key: key: Union[str, bytes], type_: str = "CONTENT", permissions: Optional[List[str]] = None, - ): + ) -> None: if isinstance(kid, str): clean_kid = kid.replace("-", "") if len(clean_kid) == 32: @@ -127,8 +127,8 @@ class CustomRemoteCDM: caching: Optional[Dict[str, Any]] = None, legacy: Optional[Dict[str, Any]] = None, timeout: int = 30, - **kwargs, - ): + **kwargs: Any, + ) -> None: """ Initialise Custom Remote CDM with highly configurable options. @@ -212,7 +212,7 @@ class CustomRemoteCDM: self.legacy_config = legacy or {"enabled": False} self._sessions: Dict[bytes, Dict[str, Any]] = {} - self._pssh_b64 = None + self._pssh_b64: Optional[str] = None self._required_kids: Optional[List[str]] = None self._http_session = Session() @@ -224,7 +224,7 @@ class CustomRemoteCDM: if custom_headers: self._http_session.headers.update(custom_headers) - def get_device_type_enum(self, device_type: str): + def get_device_type_enum(self, device_type: str) -> DeviceTypes: """Convert device type string to enum for compatibility.""" device_type_upper = device_type.upper() if device_type_upper == "ANDROID": @@ -802,7 +802,7 @@ class CustomRemoteCDM: required_kids_list = session.get("required_kids") or self._required_kids - if self.use_vaults and required_kids_list: + if self.use_vaults and self.vaults is not None and required_kids_list: vault_keys = [] for kid_str in required_kids_list: try: @@ -872,10 +872,10 @@ class CustomRemoteCDM: session["tried_cache"] = True if required_kids_list: - available_kids = set() - for key in all_available_keys: - if isinstance(key, dict) and "kid" in key: - available_kids.add(key["kid"].replace("-", "").lower()) + available_kids: set[str] = set() + for cached_key in all_available_keys: + if isinstance(cached_key, dict) and "kid" in cached_key: + available_kids.add(cached_key["kid"].replace("-", "").lower()) required_kids = set(required_kids_list) missing_kids = required_kids - available_kids @@ -896,10 +896,11 @@ class CustomRemoteCDM: if challenge and remote_session_id: if isinstance(challenge, str): + challenge_text = challenge try: - challenge = base64.b64decode(challenge) + challenge = base64.b64decode(challenge_text) except ValueError: - challenge = challenge.encode("utf-8") + challenge = challenge_text.encode("utf-8") session["challenge"] = challenge session["remote_session_id"] = remote_session_id @@ -943,10 +944,11 @@ class CustomRemoteCDM: if self.is_playready and license_message.lstrip().startswith("<"): license_message = license_message.encode("utf-8") else: + license_text = license_message try: - license_message = base64.b64decode(license_message) + license_message = base64.b64decode(license_text) except ValueError: - license_message = license_message.encode("utf-8") + license_message = license_text.encode("utf-8") pssh = session["pssh"] init_data = self.get_init_data_from_pssh(pssh, session.get("pssh_b64")) @@ -1008,8 +1010,8 @@ class CustomRemoteCDM: session.pop("cached_keys", None) session.pop("vault_keys", None) - if self.use_vaults and session["keys"]: - key_dict = {} + if self.use_vaults and self.vaults is not None and session["keys"]: + key_dict: dict[UUID | str, str] = {} for key in session["keys"]: if key["type"] == "CONTENT": try: diff --git a/unshackle/core/cdm/decrypt_labs_remote_cdm.py b/unshackle/core/cdm/decrypt_labs_remote_cdm.py index 7142d4ef..39268554 100644 --- a/unshackle/core/cdm/decrypt_labs_remote_cdm.py +++ b/unshackle/core/cdm/decrypt_labs_remote_cdm.py @@ -17,7 +17,7 @@ from unshackle.core.vaults import Vaults class MockCertificateChain: """Mock certificate chain for PlayReady compatibility.""" - def __init__(self, name: str): + def __init__(self, name: str) -> None: self._name = name def get_name(self) -> str: @@ -27,7 +27,7 @@ class MockCertificateChain: class Key: """`Key` object compatible with pywidevine.""" - def __init__(self, kid: str, key: str, type_: str = "CONTENT"): + def __init__(self, kid: str, key: str, type_: str = "CONTENT") -> None: if isinstance(kid, str): clean_kid = kid.replace("-", "") if len(clean_kid) == 32: @@ -106,8 +106,8 @@ class DecryptLabsRemoteCDM: device_type: Optional[str] = None, system_id: Optional[int] = None, security_level: Optional[int] = None, - **kwargs, - ): + **kwargs: Any, + ) -> None: """ Initialise Decrypt Labs Remote CDM for Widevine and PlayReady schemes. @@ -146,7 +146,7 @@ class DecryptLabsRemoteCDM: self.security_level = security_level or 3 self._sessions: Dict[bytes, Dict[str, Any]] = {} - self._pssh_b64 = None + self._pssh_b64: Optional[str] = None self._required_kids: Optional[List[str]] = None self._http_session = Session() self._http_session.headers.update( @@ -157,7 +157,7 @@ class DecryptLabsRemoteCDM: } ) - def get_device_type_enum(self, device_type: str): + def get_device_type_enum(self, device_type: str) -> DeviceTypes: """Convert device type string to enum for compatibility.""" device_type_upper = device_type.upper() if device_type_upper == "ANDROID": @@ -483,10 +483,10 @@ class DecryptLabsRemoteCDM: session["tried_cache"] = True if required_kids_list: - available_kids = set() - for key in all_available_keys: - if isinstance(key, dict) and "kid" in key: - available_kids.add(key["kid"].replace("-", "").lower()) + available_kids: set[str] = set() + for cached_key in all_available_keys: + if isinstance(cached_key, dict) and "kid" in cached_key: + available_kids.add(cached_key["kid"].replace("-", "").lower()) required_kids = set(required_kids_list) missing_kids = required_kids - available_kids @@ -588,10 +588,11 @@ class DecryptLabsRemoteCDM: if self.is_playready and license_message.lstrip().startswith("<"): license_message = license_message.encode("utf-8") else: + license_text = license_message try: - license_message = base64.b64decode(license_message) + license_message = base64.b64decode(license_text) except ValueError: - license_message = license_message.encode("utf-8") + license_message = license_text.encode("utf-8") pssh = session["pssh"] init_data = self.get_init_data_from_pssh(pssh, session.get("pssh_b64")) @@ -667,7 +668,7 @@ class DecryptLabsRemoteCDM: session.pop("vault_keys", None) if self.vaults and session["keys"]: - key_dict = {} + key_dict: dict[UUID | str, str] = {} for key in session["keys"]: if key["type"] == "CONTENT": try: diff --git a/unshackle/core/cdm/detect.py b/unshackle/core/cdm/detect.py index 921847ad..6365aa3c 100644 --- a/unshackle/core/cdm/detect.py +++ b/unshackle/core/cdm/detect.py @@ -48,7 +48,7 @@ def is_remote_cdm(cdm: Any) -> bool: try: from pywidevine.remotecdm import RemoteCdm as WidevineRemoteCdm except Exception: - WidevineRemoteCdm = None + WidevineRemoteCdm = None # type: ignore[misc,assignment] if WidevineRemoteCdm is not None: try: @@ -197,7 +197,7 @@ def is_widevine_cdm(cdm: Any) -> bool: try: from pywidevine.cdm import Cdm as WidevineCdm except Exception: - WidevineCdm = None + WidevineCdm = None # type: ignore[misc,assignment] if WidevineCdm is not None: try: @@ -208,7 +208,7 @@ def is_widevine_cdm(cdm: Any) -> bool: try: from pywidevine.remotecdm import RemoteCdm as WidevineRemoteCdm except Exception: - WidevineRemoteCdm = None + WidevineRemoteCdm = None # type: ignore[misc,assignment] if WidevineRemoteCdm is not None: try: diff --git a/unshackle/core/cdm/monalisa/monalisa_cdm.py b/unshackle/core/cdm/monalisa/monalisa_cdm.py index e1561759..181a9ea3 100644 --- a/unshackle/core/cdm/monalisa/monalisa_cdm.py +++ b/unshackle/core/cdm/monalisa/monalisa_cdm.py @@ -14,7 +14,7 @@ import re import sys import uuid from pathlib import Path -from typing import Dict, Union +from typing import Any, Dict, Union import wasmtime @@ -46,7 +46,7 @@ class MonaLisaCDM: "_=./this.program", ) - def __init__(self, device_path: Path): + def __init__(self, device_path: Path) -> None: """ Initialise the MonaLisa CDM. @@ -85,10 +85,10 @@ class MonaLisaCDM: except Exception as e: raise RuntimeError(f"Failed to load WASM module: {e}") - self.store = None - self.memory = None - self.instance = None - self.exports = {} + self.store: wasmtime.Store | None = None + self.memory: wasmtime.Memory | None = None + self.instance: wasmtime.Instance | None = None + self.exports: dict[str, Any] = {} self.ctx = None def open(self) -> int: @@ -225,6 +225,7 @@ class MonaLisaCDM: def extract_license_key_bytes(self) -> bytes: """Extract the 16-byte content key from WASM memory.""" + assert self.memory is not None and self.store is not None data_ptr = self.memory.data_ptr(self.store) data_len = self.memory.data_len(self.store) @@ -237,7 +238,7 @@ class MonaLisaCDM: return bytes(mem_ptr.contents[start:end]) - def ccall(self, func_name: str, return_type: type, *args): + def ccall(self, func_name: str, return_type: type, *args: Any) -> Any: """Call a WASM function with automatic string conversion.""" stack = 0 converted_args = [] @@ -274,6 +275,7 @@ class MonaLisaCDM: """Write a 32-bit integer to WASM memory.""" if addr % 4 != 0: raise ValueError(f"Unaligned i32 write: addr={addr} (must be 4-byte aligned)") + assert self.memory is not None and self.store is not None data_len = self.memory.data_len(self.store) if addr < 0 or addr + 4 > data_len: @@ -287,6 +289,7 @@ class MonaLisaCDM: """Convert string to UTF-8 and write to WASM memory.""" encoded = data.encode("utf-8") write_length = min(len(encoded), max_length - 1) + assert self.memory is not None and self.store is not None mem_data = self.memory.data_ptr(self.store) mem_ptr = ctypes.cast(mem_data, ctypes.POINTER(ctypes.c_ubyte)) @@ -298,6 +301,7 @@ class MonaLisaCDM: def write_ascii_to_memory(self, string: str, buffer: int, dont_add_null: int = 0) -> None: """Write ASCII string to WASM memory.""" + assert self.memory is not None and self.store is not None mem_data = self.memory.data_ptr(self.store) mem_ptr = ctypes.cast(mem_data, ctypes.POINTER(ctypes.c_ubyte)) @@ -308,55 +312,58 @@ class MonaLisaCDM: if dont_add_null == 0: mem_ptr[buffer + len(encoded)] = 0 - def build_imports(self): + def build_imports(self) -> list[wasmtime.Func | wasmtime.Memory]: """Assemble the WASM import stubs the MonaLisa module needs.""" + assert self.memory is not None and self.store is not None + store, memory = self.store, self.memory - def sys_fcntl64(a, b, c): + def sys_fcntl64(a: int, b: int, c: int) -> int: return 0 - def fd_write(a, b, c, d): + def fd_write(a: int, b: int, c: int, d: int) -> int: return 0 - def fd_close(a): + def fd_close(a: int) -> int: return 0 - def sys_ioctl(a, b, c): + def sys_ioctl(a: int, b: int, c: int) -> int: return 0 - def sys_open(a, b, c): + def sys_open(a: int, b: int, c: int) -> int: return 0 - def sys_rmdir(a): + def sys_rmdir(a: int) -> int: return 0 - def sys_unlink(a): + def sys_unlink(a: int) -> int: return 0 - def clock(): + def clock() -> int: return 0 - def time(a): + def time(a: int) -> int: return 0 - def emscripten_run_script(a): + def emscripten_run_script(a: int) -> None: return None - def fd_seek(a, b, c, d, e): + def fd_seek(a: int, b: int, c: int, d: int, e: int) -> int: return 0 - def emscripten_resize_heap(a): + def emscripten_resize_heap(a: int) -> int: return 0 - def fd_read(a, b, c, d): + def fd_read(a: int, b: int, c: int, d: int) -> int: return 0 - def emscripten_run_script_string(a): + def emscripten_run_script_string(a: int) -> int: return 0 - def emscripten_run_script_int(a): + def emscripten_run_script_int(a: int) -> int: return 1 - def emscripten_memcpy_big(dest, src, num): + def emscripten_memcpy_big(dest: int, src: int, num: int | None) -> int: + assert self.memory is not None and self.store is not None mem_data = self.memory.data_ptr(self.store) data_len = self.memory.data_len(self.store) if num is None: @@ -367,7 +374,7 @@ class MonaLisaCDM: mem_ptr[dest + i] = mem_ptr[src + i] return dest - def environ_get(environ_ptr, environ_buf): + def environ_get(environ_ptr: int, environ_buf: int) -> int: buf_size = 0 for index, string in enumerate(self.ENV_STRINGS): ptr = environ_buf + buf_size @@ -376,7 +383,7 @@ class MonaLisaCDM: buf_size += len(string) + 1 return 0 - def environ_sizes_get(penviron_count, penviron_buf_size): + def environ_sizes_get(penviron_count: int, penviron_buf_size: int) -> int: self.write_i32(penviron_count, len(self.ENV_STRINGS)) buf_size = sum(len(s) + 1 for s in self.ENV_STRINGS) self.write_i32(penviron_buf_size, buf_size) @@ -385,23 +392,23 @@ class MonaLisaCDM: i32 = wasmtime.ValType.i32() return [ - wasmtime.Func(self.store, wasmtime.FuncType([i32, i32, i32], [i32]), sys_fcntl64), - wasmtime.Func(self.store, wasmtime.FuncType([i32, i32, i32, i32], [i32]), fd_write), - wasmtime.Func(self.store, wasmtime.FuncType([i32], [i32]), fd_close), - wasmtime.Func(self.store, wasmtime.FuncType([i32, i32, i32], [i32]), sys_ioctl), - wasmtime.Func(self.store, wasmtime.FuncType([i32, i32, i32], [i32]), sys_open), - wasmtime.Func(self.store, wasmtime.FuncType([i32], [i32]), sys_rmdir), - wasmtime.Func(self.store, wasmtime.FuncType([i32], [i32]), sys_unlink), - wasmtime.Func(self.store, wasmtime.FuncType([], [i32]), clock), - wasmtime.Func(self.store, wasmtime.FuncType([i32], [i32]), time), - wasmtime.Func(self.store, wasmtime.FuncType([i32], []), emscripten_run_script), - wasmtime.Func(self.store, wasmtime.FuncType([i32, i32, i32, i32, i32], [i32]), fd_seek), - wasmtime.Func(self.store, wasmtime.FuncType([i32, i32, i32], [i32]), emscripten_memcpy_big), - wasmtime.Func(self.store, wasmtime.FuncType([i32], [i32]), emscripten_resize_heap), - wasmtime.Func(self.store, wasmtime.FuncType([i32, i32], [i32]), environ_get), - wasmtime.Func(self.store, wasmtime.FuncType([i32, i32], [i32]), environ_sizes_get), - wasmtime.Func(self.store, wasmtime.FuncType([i32, i32, i32, i32], [i32]), fd_read), - wasmtime.Func(self.store, wasmtime.FuncType([i32], [i32]), emscripten_run_script_string), - wasmtime.Func(self.store, wasmtime.FuncType([i32], [i32]), emscripten_run_script_int), - self.memory, + wasmtime.Func(store, wasmtime.FuncType([i32, i32, i32], [i32]), sys_fcntl64), + wasmtime.Func(store, wasmtime.FuncType([i32, i32, i32, i32], [i32]), fd_write), + wasmtime.Func(store, wasmtime.FuncType([i32], [i32]), fd_close), + wasmtime.Func(store, wasmtime.FuncType([i32, i32, i32], [i32]), sys_ioctl), + wasmtime.Func(store, wasmtime.FuncType([i32, i32, i32], [i32]), sys_open), + wasmtime.Func(store, wasmtime.FuncType([i32], [i32]), sys_rmdir), + wasmtime.Func(store, wasmtime.FuncType([i32], [i32]), sys_unlink), + wasmtime.Func(store, wasmtime.FuncType([], [i32]), clock), + wasmtime.Func(store, wasmtime.FuncType([i32], [i32]), time), + wasmtime.Func(store, wasmtime.FuncType([i32], []), emscripten_run_script), + wasmtime.Func(store, wasmtime.FuncType([i32, i32, i32, i32, i32], [i32]), fd_seek), + wasmtime.Func(store, wasmtime.FuncType([i32, i32, i32], [i32]), emscripten_memcpy_big), + wasmtime.Func(store, wasmtime.FuncType([i32], [i32]), emscripten_resize_heap), + wasmtime.Func(store, wasmtime.FuncType([i32, i32], [i32]), environ_get), + wasmtime.Func(store, wasmtime.FuncType([i32, i32], [i32]), environ_sizes_get), + wasmtime.Func(store, wasmtime.FuncType([i32, i32, i32, i32], [i32]), fd_read), + wasmtime.Func(store, wasmtime.FuncType([i32], [i32]), emscripten_run_script_string), + wasmtime.Func(store, wasmtime.FuncType([i32], [i32]), emscripten_run_script_int), + memory, ] diff --git a/unshackle/core/commands.py b/unshackle/core/commands.py index b4cff382..f1829a55 100644 --- a/unshackle/core/commands.py +++ b/unshackle/core/commands.py @@ -14,7 +14,7 @@ COMMANDS = sorted( ) -def load_command(path: Path) -> object: +def load_command(path: Path) -> click.Command: """Load one command module, returning its stem-named attribute. Raises a concise, single-line error naming the command and the real cause so @@ -32,14 +32,14 @@ def load_command(path: Path) -> object: ) from e -def load_commands(paths: list[Path]) -> tuple[dict[str, object], list[str]]: +def load_commands(paths: list[Path]) -> tuple[dict[str, click.Command], list[str]]: """Load every command, returning the good ones plus a list of load errors. Importing this module must never raise (it runs at CLI startup, before unshackle installs Rich, so a raise here prints an ugly pre-setup traceback). Instead we collect failures and surface them once, cleanly, when the user operates the CLI. """ - modules: dict[str, object] = {} + modules: dict[str, click.Command] = {} errors: list[str] = [] for path in paths: try: diff --git a/unshackle/core/config.py b/unshackle/core/config.py index 3f3a0fcf..b6f32025 100644 --- a/unshackle/core/config.py +++ b/unshackle/core/config.py @@ -90,7 +90,7 @@ class Config: core_dir = Path(__file__).resolve().parent namespace_dir = core_dir.parent commands = namespace_dir / "commands" - services = [namespace_dir / "services"] + services: list[Path | str] = [namespace_dir / "services"] vaults = namespace_dir / "vaults" fonts = namespace_dir / "fonts" user_configs = core_dir.parent @@ -114,7 +114,9 @@ class Config: chapters = "Chapters_{title}_{random}.txt" # Directories.temp subtitle = "Subtitle_{id}_{language}.srt" # Directories.temp - def __init__(self, **kwargs: Any): + repack: bool = False + + def __init__(self, **kwargs: Any) -> None: self.dl: dict = kwargs.get("dl") or {} self.cdm: dict = kwargs.get("cdm") or {} self.chapter_fallback_name: str = kwargs.get("chapter_fallback_name") or "" diff --git a/unshackle/core/console.py b/unshackle/core/console.py index 2f0efac2..809bb8cf 100644 --- a/unshackle/core/console.py +++ b/unshackle/core/console.py @@ -8,7 +8,7 @@ import unicodedata from datetime import datetime from functools import lru_cache from types import ModuleType -from typing import IO, Any, Callable, Iterable, List, Literal, Mapping, Optional, TextIO, Union +from typing import IO, Any, Callable, Iterable, List, Literal, Mapping, Optional, Sized, TextIO, Union, cast import rich.cells from rich import box @@ -213,7 +213,7 @@ class ComfyConsole(Console): get_time: Optional[Callable[[], float]] = None, _environ: Optional[Mapping[str, str]] = None, log_renderer: Optional[LogRender] = None, - ): + ) -> None: super().__init__( color_system=color_system, force_terminal=force_terminal, @@ -247,12 +247,12 @@ class ComfyConsole(Console): if log_renderer: self._log_render = log_renderer - def status( + def status( # type: ignore[override] self, status: RenderableType, *, spinner: str = "dots", - spinner_style: str = "status.spinner", + spinner_style: StyleType = "status.spinner", speed: float = 1.0, refresh_per_second: float = 12.5, pad: PaddingDimensions = (0, 5), @@ -282,8 +282,8 @@ class ComfyConsole(Console): if pad: top, right, bottom, left = Padding.unpack(pad) - renderable_width = len(status_renderable.status) - spinner_width = len(status_renderable.renderable.text) + renderable_width = len(cast(Sized, status_renderable.status)) + spinner_width = len(cast(Sized, status_renderable.renderable.text)) status_width = spinner_width + renderable_width available_width = self.width - status_width @@ -327,8 +327,8 @@ class ComfyConsole(Console): live.start() -primary_scheme = resolve_palette(config.theme) -if primary_scheme is None: +primary_scheme: dict[str, str] = resolve_palette(config.theme) or {} +if not primary_scheme: logging.getLogger("console").warning( "Unknown theme %r, falling back to %s (available: %s)", config.theme, @@ -431,9 +431,9 @@ _JOINING_VIRAMAS = frozenset("्্્୍్്") _rich_cell_len = rich.cells._cell_len -def _cell_len(text: str, *args: str) -> int: +def _cell_len(text: str, unicode_version: str) -> int: """Count a consonant joined by a virama (the न्द in हिन्दी) as no cell: terminals draw the conjunct in one.""" - width = _rich_cell_len(text, *args) + width = _rich_cell_len(text, unicode_version) if _JOINING_VIRAMAS.isdisjoint(text): return width return width - sum( diff --git a/unshackle/core/constants.py b/unshackle/core/constants.py index 4083665d..5fd93264 100644 --- a/unshackle/core/constants.py +++ b/unshackle/core/constants.py @@ -1,5 +1,8 @@ from threading import Event -from typing import TypeVar, Union +from typing import TYPE_CHECKING, TypeVar, Union + +if TYPE_CHECKING: + from unshackle.core.tracks import Audio, Subtitle, Track, Video DOWNLOAD_CANCELLED = Event() DOWNLOAD_LICENCE_ONLY = Event() @@ -46,9 +49,9 @@ context_settings = dict( # For use in signatures of functions which take one specific type of track at a time # (it can't be a list that contains e.g. both Video and Audio objects) -TrackT = TypeVar("TrackT", bound="Track") # noqa: F821 +TrackT = TypeVar("TrackT", bound="Track") # For general use in lists that can contain mixed types of tracks. # list[Track] won't work because list is invariant. # TODO: Add Chapter? -AnyTrack = Union["Video", "Audio", "Subtitle"] # noqa: F821 +AnyTrack = Union["Video", "Audio", "Subtitle"] diff --git a/unshackle/core/credential.py b/unshackle/core/credential.py index 1510d2b6..08f769e6 100644 --- a/unshackle/core/credential.py +++ b/unshackle/core/credential.py @@ -10,7 +10,7 @@ from typing import Optional, Union class Credential: """Username (or Email) and Password Credential.""" - def __init__(self, username: str, password: str, extra: Optional[str] = None): + def __init__(self, username: str, password: str, extra: Optional[str] = None) -> None: self.username = username self.password = password self.extra = extra diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index 0951546a..8f09420b 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -404,7 +404,7 @@ def no_retry_session(session: Session) -> Session: view.adapters = OrderedDict() for prefix, adapter in session.adapters.items(): twin = copy(adapter) - twin.max_retries = Retry(0, read=False) + twin.max_retries = Retry(0, read=False) # type: ignore[attr-defined] if isinstance(adapter, HTTPAdapter) and isinstance(twin, HTTPAdapter): twin.poolmanager = adapter.poolmanager twin.proxy_manager = adapter.proxy_manager @@ -869,7 +869,7 @@ def download( req_headers["Range"] = f"bytes={resume_offset}-" request_kwargs["headers"] = req_headers - stream = session.get(url, stream=True, **request_kwargs) + stream: Any = session.get(url, stream=True, **request_kwargs) if (not part_mode) and (not item_range) and resume_offset > 0 and stream.status_code == 416: complete = parse_content_range_unsatisfied( @@ -1030,7 +1030,7 @@ def download( # (unsatisfiable Range on retry, silent zero-padded corruption). part_mode still # pre-truncates in dispatch_parts because its workers seek to fixed offsets. with open(save_path if part_mode else tmp_file, file_mode, buffering=file_buffering) as f: - if part_mode: + if part_mode and part_offset is not None: # spelled so mypy narrows f.seek(part_offset + written) _write = f.write @@ -1493,7 +1493,7 @@ def download_multiprocess( def requests( - urls: Union[str, list[str], dict[str, Any], list[dict[str, Any]]], + urls: Union[str, dict[str, Any], list[str], list[dict[str, Any]], list[Union[str, dict[str, Any]]]], output_dir: Path, filename: str, headers: Optional[MutableMapping[str, Union[str, bytes]]] = None, @@ -2106,16 +2106,16 @@ def requests( completed, pending = wait(pending, timeout=0.1, return_when=FIRST_COMPLETED) for future in completed: - exc = future.exception() - if isinstance(exc, KeyboardInterrupt): + err = future.exception() + if isinstance(err, KeyboardInterrupt): raise KeyboardInterrupt() - elif isinstance(exc, TailPartFailed) and permanent_status(exc) is None: - if exc.index not in tail_skipped: - tail_boosted.discard(exc.index) - tail_skipped.add(exc.index) - remaining.appendleft((exc.index, urls[exc.index])) + elif isinstance(err, TailPartFailed) and permanent_status(err) is None: + if err.index not in tail_skipped: + tail_boosted.discard(err.index) + tail_skipped.add(err.index) + remaining.appendleft((err.index, urls[err.index])) continue - elif exc: + elif err: DOWNLOAD_CANCELLED.set() yield dict(downloaded="[red]FAILING") pool.shutdown(wait=False, cancel_futures=True) @@ -2124,14 +2124,14 @@ def requests( debug_logger.log( level="ERROR", operation="downloader_failed", - message=f"Download failed: {exc}", - error=exc, + message=f"Download failed: {err}", + error=err, context={ "url_count": len(urls), "output_dir": str(output_dir), }, ) - raise exc + raise err except KeyboardInterrupt: DOWNLOAD_CANCELLED.set() yield dict(downloaded="[yellow]CANCELLING") diff --git a/unshackle/core/drm/blowfish_stripe.py b/unshackle/core/drm/blowfish_stripe.py index 29530f0a..8e7d3b9c 100644 --- a/unshackle/core/drm/blowfish_stripe.py +++ b/unshackle/core/drm/blowfish_stripe.py @@ -17,7 +17,7 @@ class BlowfishStripe: CHUNK = 2048 IV = bytes(range(8)) - def __init__(self, key: Union[bytes, str]): + def __init__(self, key: Union[bytes, str]) -> None: self.key: bytes = bytes.fromhex(key) if isinstance(key, str) else key def decrypt_chunk(self, chunk: bytes) -> bytes: diff --git a/unshackle/core/drm/clearkey.py b/unshackle/core/drm/clearkey.py index 856bb568..dcfbbf12 100644 --- a/unshackle/core/drm/clearkey.py +++ b/unshackle/core/drm/clearkey.py @@ -12,7 +12,7 @@ from Cryptodome.Util.Padding import unpad from m3u8.model import Key from requests import Session -from unshackle.core.session import RnetSession +from unshackle.core.session import AnySession, RnetSession from unshackle.core.utilities import log_event from unshackle.core.utils import run_proxy @@ -28,7 +28,7 @@ TS_SYNC_CHECKS = 4 class ClearKey: """AES-128 ClearKey (HLS) DRM system.""" - def __init__(self, key: Union[bytes, str], iv: Optional[Union[bytes, str]] = None): + def __init__(self, key: Union[bytes, str], iv: Optional[Union[bytes, str]] = None) -> None: """ Give the IV where possible. If you do not give an IV, unshackle sets it to \x00 of the same bit-size as the content key. @@ -104,7 +104,7 @@ class ClearKey: shutil.move(decrypted_path, path) @classmethod - def from_m3u_key(cls, m3u_key: Key, session: Optional[Session] = None) -> ClearKey: + def from_m3u_key(cls, m3u_key: Key, session: Optional[AnySession] = None) -> ClearKey: """ Load a ClearKey from an M3U(8) Playlist's EXT-X-KEY. diff --git a/unshackle/core/drm/clearkey_cenc.py b/unshackle/core/drm/clearkey_cenc.py index 0bdc699d..c62a0726 100644 --- a/unshackle/core/drm/clearkey_cenc.py +++ b/unshackle/core/drm/clearkey_cenc.py @@ -45,7 +45,7 @@ class ClearKeyCENC: laurl: Optional[str] = None, content_keys: Optional[dict[UUID, str]] = None, **kwargs: Any, - ): + ) -> None: kid_list: list[UUID] = [] for kid in kids or []: if isinstance(kid, str): @@ -279,7 +279,7 @@ class ClearKeyCENC: config.directories.temp.mkdir(parents=True, exist_ok=True) try: - arguments = [ + arguments: list[str | Path] = [ f"input={path},stream=0,output={output_path},output_format=MP4", "--enable_raw_key_decryption", "--keys", @@ -304,6 +304,7 @@ class ClearKeyCENC: had_error = False shaka_log_buffer = "" + assert p.stderr is not None for line in iter(p.stderr.readline, ""): line = line.strip() if not line: diff --git a/unshackle/core/drm/monalisa.py b/unshackle/core/drm/monalisa.py index e0c92838..32a96bd9 100644 --- a/unshackle/core/drm/monalisa.py +++ b/unshackle/core/drm/monalisa.py @@ -41,7 +41,7 @@ class MonaLisa: device_path: Path, tool_path: Optional[Path] = None, **kwargs: Any, - ): + ) -> None: """ Initialise MonaLisa DRM. diff --git a/unshackle/core/drm/playready.py b/unshackle/core/drm/playready.py index 6a30d4a3..0e97adae 100644 --- a/unshackle/core/drm/playready.py +++ b/unshackle/core/drm/playready.py @@ -18,7 +18,6 @@ from pyplayready.cdm import Cdm as PlayReadyCdm from pyplayready.misc.exceptions import InvalidPssh from pyplayready.system.pssh import PSSH from pyplayready.system.wrmheader import WRMHeader -from requests import Session from rich.text import Text from unshackle.core import binaries @@ -26,7 +25,7 @@ from unshackle.core.config import config from unshackle.core.console import console from unshackle.core.constants import AnyTrack from unshackle.core.drm import key_args as drm_key_args -from unshackle.core.session import utf8_text +from unshackle.core.session import AnySession, utf8_text from unshackle.core.utilities import get_boxes, log_event from unshackle.core.utils import run_proxy from unshackle.core.utils.subprocess import ffprobe @@ -43,7 +42,7 @@ class PlayReady: kid: Union[UUID, str, bytes, None] = None, pssh_b64: Optional[str] = None, **kwargs: Any, - ): + ) -> None: if not pssh: raise ValueError("Provided PSSH is empty.") if not isinstance(pssh, PSSH): @@ -192,7 +191,7 @@ class PlayReady: return cls(pssh=pssh, kid=kid, pssh_b64=pssh_b64) @classmethod - def from_track(cls, track: AnyTrack, session: Optional[Session] = None) -> PlayReady: + def from_track(cls, track: AnyTrack, session: Optional[AnySession] = None) -> PlayReady: if not session: session = run_proxy.fallback_session() session.headers.update(config.headers) @@ -203,6 +202,7 @@ class PlayReady: if track.descriptor == track.Descriptor.HLS: m3u_url = track.url + assert isinstance(m3u_url, str) master = m3u8.loads(utf8_text(session.get(m3u_url)), uri=m3u_url) pssh_boxes.extend( Box.parse(base64.b64decode(x.uri.split(",")[-1])) @@ -651,7 +651,7 @@ class PlayReady: config.directories.temp.mkdir(parents=True, exist_ok=True) try: - arguments = [ + arguments: list[str | Path] = [ f"input={path},stream=0,output={output_path},output_format=MP4", "--enable_raw_key_decryption", "--keys", @@ -672,6 +672,7 @@ class PlayReady: stream_skipped = False had_error = False shaka_log_buffer = "" + assert p.stderr is not None for line in iter(p.stderr.readline, ""): line = line.strip() if not line: diff --git a/unshackle/core/drm/playready_errors.py b/unshackle/core/drm/playready_errors.py index f58452da..c29ded1e 100644 --- a/unshackle/core/drm/playready_errors.py +++ b/unshackle/core/drm/playready_errors.py @@ -873,7 +873,7 @@ def describe(code: int) -> str | None: if __name__ == "__main__": assert describe(0x8004C065) == describe(-2147172251) - assert "DEVCERT_REVOKED" in describe(-2147172251) + assert "DEVCERT_REVOKED" in (describe(-2147172251) or "") assert describe(-1021) is None # non-DRM_E service code assert is_revocation(0x8004C065) and is_revocation(-2147172251) # DEVCERT_REVOKED, both forms assert not is_revocation(0x8004C006) # INVALIDLICENSE is not revocation diff --git a/unshackle/core/drm/widevine.py b/unshackle/core/drm/widevine.py index 5f3039fe..d3fac0bd 100644 --- a/unshackle/core/drm/widevine.py +++ b/unshackle/core/drm/widevine.py @@ -14,7 +14,6 @@ from construct import Container from pymp4.parser import Box from pywidevine.cdm import Cdm as WidevineCdm from pywidevine.pssh import PSSH -from requests import Session from rich.text import Text from unshackle.core import binaries @@ -22,7 +21,7 @@ from unshackle.core.config import config from unshackle.core.console import console from unshackle.core.constants import AnyTrack from unshackle.core.drm import key_args as drm_key_args -from unshackle.core.session import utf8_text +from unshackle.core.session import AnySession, utf8_text from unshackle.core.utilities import get_boxes, log_event from unshackle.core.utils import run_proxy from unshackle.core.utils.subprocess import ffprobe @@ -37,7 +36,7 @@ class Widevine: UUID("00010203-0405-0607-0809-101112131415"), # Shaka Packager test pattern } - def __init__(self, pssh: PSSH, kid: Union[UUID, str, bytes, None] = None, **kwargs: Any): + def __init__(self, pssh: PSSH, kid: Union[UUID, str, bytes, None] = None, **kwargs: Any) -> None: if not pssh: raise ValueError("Provided PSSH is empty.") if not isinstance(pssh, PSSH): @@ -69,7 +68,7 @@ class Widevine: self.data: dict = kwargs or {} @classmethod - def from_track(cls, track: AnyTrack, session: Optional[Session] = None) -> Widevine: + def from_track(cls, track: AnyTrack, session: Optional[AnySession] = None) -> Widevine: """ Get PSSH and KID from within the Initiation Segment of the Track Data. It also tries to get PSSH and KID from other track data like M3U8 data @@ -96,6 +95,7 @@ class Widevine: if track.descriptor == track.Descriptor.HLS: m3u_url = track.url + assert isinstance(m3u_url, str) master = m3u8.loads(utf8_text(session.get(m3u_url)), uri=m3u_url) pssh_boxes.extend( Box.parse(base64.b64decode(x.uri.split(",")[-1])) @@ -380,7 +380,7 @@ class Widevine: config.directories.temp.mkdir(parents=True, exist_ok=True) try: - arguments = [ + arguments: list[str | Path] = [ f"input={path},stream=0,output={output_path},output_format=MP4", "--enable_raw_key_decryption", "--keys", @@ -402,6 +402,7 @@ class Widevine: had_error = False shaka_log_buffer = "" + assert p.stderr is not None for line in iter(p.stderr.readline, ""): line = line.strip() if not line: diff --git a/unshackle/core/events.py b/unshackle/core/events.py index a85451ef..fe09ff36 100644 --- a/unshackle/core/events.py +++ b/unshackle/core/events.py @@ -15,12 +15,12 @@ class Events: # Track is about to be Multiplexed into a Container TRACK_MULTIPLEX = 5 - def __init__(self): + def __init__(self) -> None: self.__subscriptions: dict[Events.Types, list[Callable]] = {} self.__ephemeral: dict[Events.Types, list[Callable]] = {} self.reset() - def reset(self): + def reset(self) -> None: """Reset Event Observer clearing all Subscriptions.""" self.__subscriptions = {k: [] for k in Events.Types.__members__.values()} self.__ephemeral = deepcopy(self.__subscriptions) diff --git a/unshackle/core/import_service.py b/unshackle/core/import_service.py index ecbbab42..0d78b58e 100644 --- a/unshackle/core/import_service.py +++ b/unshackle/core/import_service.py @@ -16,7 +16,7 @@ from langcodes import tag_is_valid from unshackle.core.config import config from unshackle.core.constants import AnyTrack from unshackle.core.credential import Credential -from unshackle.core.drm import drm_from_dict, real_kids +from unshackle.core.drm import ClearKey, MonaLisa, drm_from_dict, real_kids from unshackle.core.manifests import DASH, HLS, ISM, MissingLanguageError from unshackle.core.remote_service import ( RemoteService, @@ -35,7 +35,7 @@ from unshackle.core.utils import run_proxy log = logging.getLogger("import") E = TypeVar("E", bound=Enum) -PARSERS = {"DASH": DASH, "HLS": HLS, "ISM": ISM} +PARSERS: dict[str, type[DASH] | type[HLS] | type[ISM]] = {"DASH": DASH, "HLS": HLS, "ISM": ISM} MANIFEST_DATA_KEYS = {"DASH": "dash", "ISM": "ism"} @@ -399,8 +399,8 @@ class ImportService: ) except Exception as e: raise click.ClickException(f"Failed to parse the {manifest_type} manifest for '{title}': {e}") - for track in parsed: - tracks.add(track) + for parsed_track in parsed: + tracks.add(parsed_track) parsed_ids = {str(t.id) for t in tracks} url_dicts = [t for t in tracks_map.values() if t.get("descriptor", "URL") == "URL"] # Merge back only tracks from another manifest, such as a separate HEVC one; the @@ -666,7 +666,10 @@ class ImportService: continue keys = {kid: key for kid, key in pool.items() if kid in declared} if declared else pool if track.drm: - drm_obj = track.drm[0] + drm_obj: Any = next(iter(track.drm)) + # HLS AES-128 fetches its own key, and MonaLisa's content_keys is a read-only view. + if isinstance(drm_obj, (ClearKey, MonaLisa)): + continue else: drm_obj = RemoteService.create_drm_stub(system, [kid.hex for kid in (declared or pool)]) for kid, key in keys.items(): diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index b904e449..c4630811 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -11,7 +11,7 @@ from contextlib import ExitStack, closing from copy import deepcopy from functools import lru_cache, partial from pathlib import Path -from typing import Any, BinaryIO, Callable, Optional, Union +from typing import Any, BinaryIO, Callable, Optional, Union, cast from urllib.parse import urljoin, urlparse from uuid import UUID from zlib import crc32 @@ -29,13 +29,14 @@ from unshackle.core.config import config from unshackle.core.constants import DOWNLOAD_ALL_DRM, DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack from unshackle.core.downloaders.requests import parse_content_range from unshackle.core.drm import DRM_T, ClearKeyCENC, PlayReady, Widevine -from unshackle.core.drm.segment_decrypt import SegmentDecrypter, can_use +from unshackle.core.drm.segment_decrypt import SEGMENT_DRM_T, SegmentDecrypter, can_use from unshackle.core.drm.verify import decrypt_track from unshackle.core.events import events from unshackle.core.session import RnetSession from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video, resume from unshackle.core.tracks.track import ( DRM_PREFERENCE_TYPES, + Track, assert_fragments_decrypted, parse_sidx, ) @@ -117,7 +118,7 @@ class RollingMerge: class DASH: - def __init__(self, manifest, url: str): + def __init__(self, manifest: Element, url: str) -> None: if manifest is None: raise ValueError("DASH manifest must be provided.") if manifest.tag != "MPD": @@ -149,7 +150,7 @@ class DASH: url = res.url if not res.ok: - raise requests.ConnectionError("Failed to request the MPD document.", response=res) + raise requests.ConnectionError("Failed to request the MPD document.", response=res) # type: ignore[arg-type] log_event( "manifest_dash_fetch", @@ -237,7 +238,7 @@ class DASH: # likely mp4-boxed subtitles # TODO: It may not actually be subtitles try: - real_codec = Subtitle.Codec.from_mime(codecs) + real_codec = Subtitle.Codec.from_mime(codecs or "") content_type = "text" mime_type = f"application/mp4; codecs='{real_codec.value.lower()}'" except ValueError: @@ -248,9 +249,12 @@ class DASH: codecs = mime_type.split("/")[1] video_args: dict[str, Any] = {} + track_type: type[Video] | type[Audio] | type[Subtitle] + track_codec: Any # always the Codec of track_type, a pairing mypy cannot express + track_args: dict[str, Any] if content_type == "video": track_type = Video - track_codec = Video.Codec.from_codecs(codecs) + track_codec = Video.Codec.from_codecs(codecs or "") track_fps = get("frameRate") if not track_fps and segment_base is not None: track_fps = segment_base.get("timescale") @@ -262,7 +266,7 @@ class DASH: track_args = dict( range_=self.get_video_range( - codecs, findall("SupplementalProperty"), findall("EssentialProperty") + codecs or "", findall("SupplementalProperty"), findall("EssentialProperty") ), bitrate=get("bandwidth") or None, width=get("width") or 0, @@ -275,7 +279,7 @@ class DASH: ) elif content_type == "audio": track_type = Audio - track_codec = Audio.Codec.from_codecs(codecs) + track_codec = Audio.Codec.from_codecs(codecs or "") track_args = dict( bitrate=get("bandwidth") or None, channels=next( @@ -384,6 +388,7 @@ class DASH: session.proxies.update({"all": proxy}) log = logging.getLogger("DASH") + assert isinstance(track.url, str) manifest: ElementTree = track.data["dash"]["manifest"] adaptation_set: Element = track.data["dash"]["adaptation_set"] @@ -412,7 +417,7 @@ class DASH: or (existing_drm and manifest_drm and not any(isinstance(drm, Widevine) for drm in existing_drm)) ) - pre_existing_keys = {} + pre_existing_keys: dict[UUID, str] = {} if existing_drm: for drm_obj in existing_drm: if hasattr(drm_obj, "content_keys") and drm_obj.content_keys: @@ -421,6 +426,7 @@ class DASH: if should_override_drm: track.drm = manifest_drm else: + assert existing_drm is not None track.drm = existing_drm if DOWNLOAD_ALL_DRM.is_set(): track.drm = [ @@ -592,6 +598,7 @@ class DASH: if can_resume: resume.write_sidecar(save_dir, digest) + downloader_args: dict[str, Any] if collapse_single_url: # must carry no Range header, or the downloader would skip its ranged-parallel path downloader_args = dict( @@ -644,7 +651,7 @@ class DASH: decrypter: Optional[SegmentDecrypter] = None if use_segment_decrypt and init_data: - decrypter = SegmentDecrypter(drm, init_data, save_dir.parent, max_workers) + decrypter = SegmentDecrypter(cast(SEGMENT_DRM_T, drm), init_data, save_dir.parent, max_workers) merger: Optional[RollingMerge] = None try: @@ -796,7 +803,7 @@ class DASH: return element if element is not None else adaptation_set.find(name) @staticmethod - def has_segment_timing(track: AnyTrack) -> bool: + def has_segment_timing(track: Track) -> bool: """Return False when the manifest gives the track as one file with no index, so a sample cannot cut it.""" dash = track.data.get("dash") or {} representation, adaptation_set = dash.get("representation"), dash.get("adaptation_set") @@ -921,7 +928,7 @@ class DASH: adaptation_set: Element, representation: Element, manifest: ElementTree, - track: AnyTrack, + track: Track, track_url: str, session: Union[Session, RnetSession], probe_kid: bool = True, @@ -1225,6 +1232,7 @@ class DASH: if not tag_is_valid(option) or option.startswith("und"): continue return Language.get(option) + return None @staticmethod def get_video_range( diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index 8cdab1bb..fdde9b0d 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -12,7 +12,7 @@ import time from contextlib import ExitStack, closing from itertools import accumulate from pathlib import Path -from typing import Any, Optional, Union +from typing import Any, Optional, Union, cast from urllib.parse import urljoin, urlparse from uuid import UUID from zlib import crc32 @@ -31,7 +31,7 @@ from unshackle.core import binaries from unshackle.core.cdm.detect import cdm_type_stub, is_playready_cdm, is_widevine_cdm from unshackle.core.config import config from unshackle.core.constants import DOWNLOAD_ALL_DRM, DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack -from unshackle.core.drm import DRM_T, ClearKey, MonaLisa, PlayReady, Widevine +from unshackle.core.drm import DRM_T, ClearKey, ClearKeyCENC, MonaLisa, PlayReady, Widevine from unshackle.core.drm.segment_decrypt import SegmentDecrypter, can_use from unshackle.core.drm.verify import decrypt_track from unshackle.core.events import events @@ -54,7 +54,7 @@ class HLS: session: Optional[Union[Session, RnetSession]] = None, url: Optional[str] = None, raw_text: Optional[str] = None, - ): + ) -> None: if not manifest: raise ValueError("HLS manifest must be provided.") if not isinstance(manifest, M3U8): @@ -84,7 +84,7 @@ class HLS: if not isinstance(res, (requests.Response, RnetResponse)): raise TypeError(f"Expected response to be a requests.Response or rnet.Response, not {type(res)}") if not res.ok: - raise requests.ConnectionError("Failed to request the M3U(8) document.", response=res) + raise requests.ConnectionError("Failed to request the M3U(8) document.", response=res) # type: ignore[arg-type] content = utf8_text(res) master = m3u8.loads(content, uri=url) @@ -174,7 +174,7 @@ class HLS: ) return candidates[0][1] - def to_tracks(self, language: Union[str, Language]) -> Tracks: + def to_tracks(self, language: Optional[Union[str, Language]]) -> Tracks: """ Convert a Variant Playlist M3U(8) document to Video, Audio and Subtitle Track objects. @@ -224,6 +224,7 @@ class HLS: for playlist in self.manifest.playlists: video_codec: Optional[Video.Codec] = None + primary_track_type: type[Audio] | type[Video] try: # TODO: Any better way to figure out the primary track type? if playlist.stream_info.codecs: @@ -251,7 +252,7 @@ class HLS: id_=hex(crc32(str(playlist).encode()))[2:], url=urljoin(playlist.base_uri, playlist.uri), codec=( - primary_track_type.Codec.from_codecs(playlist.stream_info.codecs) + cast(Any, primary_track_type.Codec.from_codecs(playlist.stream_info.codecs)) if playlist.stream_info.codecs else None ), @@ -283,11 +284,10 @@ class HLS: continue joc = 0 + track_type: type[Audio] | type[Subtitle] if media.type == "AUDIO": track_type = Audio - codec: Union[Audio.Codec, Subtitle.Codec, None] = HLS.select_audio_codec( - audio_codecs_by_group_id.get(media.group_id), media.uri - ) + codec: Any = HLS.select_audio_codec(audio_codecs_by_group_id.get(media.group_id), media.uri) if media.channels: # the channel count, then /-separated coding (JOC) and usage identifiers (RFC 8216bis) count, *identifiers = media.channels.split("/") @@ -345,6 +345,7 @@ class HLS: has_codec = video.codec is not None if has_resolution and has_codec: continue + assert isinstance(video.url, str) try: probe = HLS.probe_ts_info(video.url, self.session) if probe: @@ -641,6 +642,7 @@ class HLS: session.proxies.update({"all": proxy}) log = logging.getLogger("HLS") + assert isinstance(track.url, str) if track.from_file: master = m3u8.load(str(track.from_file)) @@ -832,7 +834,7 @@ class HLS: if config.continue_downloads and resumable_drm: resume.write_sidecar(save_dir, digest) - downloader_args = dict( + downloader_args: dict[str, Any] = dict( urls=urls, output_dir=segment_save_dir, filename="{i:0%d}{ext}" % len(str(len(urls))), @@ -866,6 +868,7 @@ class HLS: init_content, _ = HLS.fetch_init_section(session, init_section, 0) map_data = (init_section, init_content) if use_segment_decrypt: + assert isinstance(session_drm, (Widevine, PlayReady, ClearKeyCENC)) segment_decrypter = SegmentDecrypter(session_drm, init_content, save_dir.parent, max_workers) merger: Optional[RollingMerge] = None @@ -957,10 +960,9 @@ class HLS: # no explicit IV (RFC 8216 §5.2), where each segment's IV is its sequence number. media_sequence_start = getattr(master, "media_sequence", None) or 0 decrypted_init: Optional[bytes] = None + encryption_data: Optional[tuple[Optional[m3u8.Key], DRM_T]] = None if session_drm: - encryption_data: Optional[tuple[Optional[m3u8.Key], DRM_T]] = (initial_drm_key, session_drm) - else: - encryption_data: Optional[tuple[Optional[m3u8.Key], DRM_T]] = None + encryption_data = (initial_drm_key, session_drm) i = -1 for real_i, segment in enumerate(master.segments): @@ -969,7 +971,7 @@ class HLS: is_last_segment = (real_i + 1) == len(master.segments) - def merge(to: Path, via: list[Path], delete: bool = False, include_map_data: bool = False): + def merge(to: Path, via: list[Path], delete: bool = False, include_map_data: bool = False) -> None: """ Merge all files to a given path, optionally including map data. @@ -988,6 +990,7 @@ class HLS: pass with open(to, "wb") as x: if prepend: + assert map_data is not None x.write(map_data[1]) for file in via: with open(file, "rb") as segment: @@ -996,7 +999,7 @@ class HLS: if delete: file.unlink() - def decrypt(include_this_segment: bool) -> Path: + def decrypt(include_this_segment: bool) -> Optional[Path]: """ Decrypt all segments that uses the currently set DRM. @@ -1014,6 +1017,7 @@ class HLS: Returns the decrypted path. """ nonlocal map_data, decrypted_init + assert encryption_data is not None drm = encryption_data[1] first_segment_i = next( (int(file.stem) for file in sorted(segment_save_dir.iterdir()) if file.stem.isdigit()), None @@ -1058,7 +1062,7 @@ class HLS: # output decrypted to garbage. seq_iv = isinstance(drm, ClearKey) and not (key_obj and key_obj.iv) for file in files: - if seq_iv and file.stem.isdigit(): + if seq_iv and isinstance(drm, ClearKey) and file.stem.isdigit(): download_i = int(file.stem) seq = getattr(wanted_segments[download_i], "media_sequence", None) if seq is None: @@ -1071,7 +1075,7 @@ class HLS: return decrypted_path - def merge_discontinuity(include_this_segment: bool, include_map_data: bool = True): + def merge_discontinuity(include_this_segment: bool, include_map_data: bool = True) -> None: """ Merge all segments of the discontinuity. @@ -1142,6 +1146,8 @@ class HLS: init_kid = track.get_key_id(map_data[1]) if map_data and key.method != "AES-128" else None drm = HLS.get_drm(key, session, **({"kid": init_kid} if init_kid else {})) if isinstance(drm, (Widevine, PlayReady)): + if not license_widevine: + raise ValueError("license_widevine func must be supplied to use DRM") track_kid = init_kid or drm.kid progress(downloaded="LICENSING") license_widevine(drm, track_kid=track_kid) @@ -1154,6 +1160,7 @@ class HLS: encryption_data and isinstance(drm, (Widevine, PlayReady)) and isinstance(encryption_data[1], type(drm)) + and isinstance(encryption_data[1], (Widevine, PlayReady)) and getattr(encryption_data[1], "content_keys", None) ): for prev_kid, prev_key in encryption_data[1].content_keys.items(): @@ -1268,7 +1275,7 @@ class HLS: break current_dir = current_dir.parent - def cleanup_segments_and_dirs(): + def cleanup_segments_and_dirs() -> None: """Clean up segments and directories after successful merge.""" for segment in segments: segment.unlink(missing_ok=True) @@ -1556,6 +1563,7 @@ class HLS: if not session: session = run_proxy.fallback_session() + drm: DRM_T # TODO: Add support for 'SAMPLE-AES', 'AES-CTR', 'AES-CBC', 'ClearKey' if key.method == "AES-128": drm = ClearKey.from_m3u_key(key, session) @@ -1625,9 +1633,14 @@ class HLS: if any(key.method == "NONE" for key in keys): return [] + session = run_proxy.fallback_session() + if proxy: + session.proxies.update({"all": proxy}) + run_proxy.pin(session) + for key in keys: try: - drm = HLS.get_drm(key, proxy) + drm = HLS.get_drm(key, session) drm_objects.append(drm) except NotImplementedError: unsupported_keys.append(key) diff --git a/unshackle/core/manifests/ism.py b/unshackle/core/manifests/ism.py index 8d10b6e9..f0f7a3a7 100644 --- a/unshackle/core/manifests/ism.py +++ b/unshackle/core/manifests/ism.py @@ -446,6 +446,7 @@ class ISM: use_rolling_merge = config.merge_segments and not isinstance(track, Subtitle) can_resume = config.continue_downloads and not use_rolling_merge + assert isinstance(track.url, str) digest = resume.fingerprint(track.url, [(url, None) for url in segments]) if not (can_resume and resume.reusable(save_dir, digest)): shutil.rmtree(save_dir, ignore_errors=True) @@ -453,7 +454,7 @@ class ISM: if can_resume: resume.write_sidecar(save_dir, digest) - downloader_args = dict( + downloader_args: dict[str, Any] = dict( urls=[{"url": url} for url in segments], output_dir=save_dir, filename="{i:0%d}.mp4" % len(str(len(segments))), diff --git a/unshackle/core/manifests/ism_init.py b/unshackle/core/manifests/ism_init.py index 971ea7e6..1cb4bf5f 100644 --- a/unshackle/core/manifests/ism_init.py +++ b/unshackle/core/manifests/ism_init.py @@ -765,7 +765,6 @@ def build_init_segment( raise ValueError(f"Unsupported stream type: {stream_type}") fourcc = (fourcc or "").upper() cpd = binascii.unhexlify(codec_private_data) if codec_private_data else b"" - encrypted = kid is not None # mdhd packs exactly three a-z letters; anything else (2-letter tags, # uppercase) would underflow the 5-bit fields, so fall back to "und". lang = (language or "").lower() @@ -876,7 +875,7 @@ def build_init_segment( else: raise NotImplementedError(f"Unsupported video FourCC: {fourcc}") sample_entry_payload += config_box - if encrypted: + if kid is not None: sample_entry_payload += build_sinf(codec_fourcc, kid, iv_size, constant_iv) sample_entry_box = box(b"encv", sample_entry_payload) else: @@ -904,7 +903,7 @@ def build_init_segment( codec_fourcc = b"ec-3" else: raise NotImplementedError(f"Unsupported audio FourCC: {fourcc}") - if encrypted: + if kid is not None: sample_entry_payload += build_sinf(codec_fourcc, kid, iv_size, constant_iv) sample_entry_box = box(b"enca", sample_entry_payload) else: diff --git a/unshackle/core/music/renderer.py b/unshackle/core/music/renderer.py index a93d7ce4..c40ed176 100644 --- a/unshackle/core/music/renderer.py +++ b/unshackle/core/music/renderer.py @@ -35,7 +35,7 @@ class MusicRenderer: first_song = music[0] data = first_song.data if isinstance(first_song.data, dict) else {} - metadata = data.get("metadata") if isinstance(data.get("metadata"), dict) else {} + metadata = m if isinstance(m := data.get("metadata"), dict) else {} title = music.title or first_song.album artist = music.artist or first_song.album_artist or first_song.artist year = getattr(music, "year", None) or first_song.year @@ -139,7 +139,7 @@ class MusicRenderer: def option_from_song(self, song: Song) -> Text: data = song.data if isinstance(song.data, dict) else {} - metadata = data.get("metadata") if isinstance(data.get("metadata"), dict) else {} + metadata = m if isinstance(m := data.get("metadata"), dict) else {} quality = first_text(data.get("quality"), metadata.get("quality"), metadata.get("quality_label")) raw_duration = self.first_value(data.get("duration"), metadata.get("duration")) @@ -304,7 +304,7 @@ class MusicRenderer: total = 0 for song in music: data = song.data if isinstance(song.data, dict) else {} - metadata = data.get("metadata") if isinstance(data.get("metadata"), dict) else {} + metadata = m if isinstance(m := data.get("metadata"), dict) else {} value = MusicRenderer.first_value(data.get("duration"), metadata.get("duration")) try: total += int(float(value or 0)) diff --git a/unshackle/core/music/tagger.py b/unshackle/core/music/tagger.py index 771dd9f8..de288aef 100644 --- a/unshackle/core/music/tagger.py +++ b/unshackle/core/music/tagger.py @@ -39,11 +39,11 @@ try: from mutagen.oggopus import OggOpus from mutagen.oggvorbis import OggVorbis except ImportError: # pragma: no cover - optional tagging dependency - FLAC = Picture = None - MP3 = MP4 = MP4Cover = OggOpus = OggVorbis = None # type: ignore[assignment] - ID3NoHeaderError = Exception - APIC = COMM = TALB = TCOM = TCON = TCOP = TDRC = TIT2 = TPE1 = TPE2 = TPOS = TRCK = TSRC = TXXX = None - USLT = None + FLAC = Picture = None # type: ignore[assignment, misc] + MP3 = MP4 = MP4Cover = OggOpus = OggVorbis = None # type: ignore[assignment, misc] + ID3NoHeaderError = Exception # type: ignore[assignment, misc] + APIC = COMM = TALB = TCOM = TCON = TCOP = TDRC = TIT2 = TPE1 = TPE2 = TPOS = TRCK = TSRC = TXXX = None # type: ignore[assignment, misc] + USLT = None # type: ignore[assignment, misc] log = logging.getLogger("MUSIC_TAGGER") @@ -202,6 +202,7 @@ def write_mp3_tags(path: Path, tags: dict[str, str], cover_data: Optional[bytes] audio.add_tags() if audio.tags is None: audio.add_tags() + assert audio.tags is not None frame_map = { "TITLE": ("TIT2", TIT2), diff --git a/unshackle/core/proxies/basic.py b/unshackle/core/proxies/basic.py index afda6cc3..7a39e300 100644 --- a/unshackle/core/proxies/basic.py +++ b/unshackle/core/proxies/basic.py @@ -10,7 +10,7 @@ from unshackle.core.utils.redact import mask_proxy class Basic(Proxy): - def __init__(self, **countries: Union[str, list[str]]): + def __init__(self, **countries: Union[str, list[str]]) -> None: """Basic proxy provider that uses the proxies given in the configuration file.""" self.countries = {k.lower(): v for k, v in countries.items()} diff --git a/unshackle/core/proxies/controld.py b/unshackle/core/proxies/controld.py index b047bd6f..a0cfc1a7 100644 --- a/unshackle/core/proxies/controld.py +++ b/unshackle/core/proxies/controld.py @@ -91,7 +91,7 @@ class IPv4Adapter(HTTPAdapter): class Resolver: - def __init__(self, url: str): + def __init__(self, url: str) -> None: """Resolve hostnames through one Control D DoH endpoint, caching answers for their TTL.""" self.url = url self.session = requests.Session() @@ -268,7 +268,7 @@ class Handler(socketserver.StreamRequestHandler): class LocalProxy(socketserver.ThreadingTCPServer): daemon_threads = True - def __init__(self, resolver: Resolver): + def __init__(self, resolver: Resolver) -> None: """An HTTP forward proxy on loopback that resolves every host through `resolver`.""" super().__init__(("127.0.0.1", 0), Handler) self.resolver = resolver @@ -331,7 +331,7 @@ class ControlD(Proxy): resolver: Optional[str] = None, profiles: Optional[list[dict]] = None, max_profiles: int = 4, - ): + ) -> None: """ Proxy provider that uses Control D's transparent proxies. diff --git a/unshackle/core/proxies/expressvpn.py b/unshackle/core/proxies/expressvpn.py index b0998ce4..73e10193 100644 --- a/unshackle/core/proxies/expressvpn.py +++ b/unshackle/core/proxies/expressvpn.py @@ -54,7 +54,7 @@ class ExpressVPN(Proxy): cache_path: Optional[str] = None, timeout: float = 10.0, enable: bool = False, - ): + ) -> None: """ Proxy provider that uses ExpressVPN app proxy credentials. diff --git a/unshackle/core/proxies/gluetun.py b/unshackle/core/proxies/gluetun.py index d58fcd63..81037db4 100644 --- a/unshackle/core/proxies/gluetun.py +++ b/unshackle/core/proxies/gluetun.py @@ -7,7 +7,7 @@ import subprocess import tempfile import threading import time -from typing import Optional +from typing import Any, Optional import requests @@ -22,7 +22,7 @@ cleanup_lock = threading.Lock() cleanup_registered = False -def cleanup_all_gluetun_containers(): +def cleanup_all_gluetun_containers() -> None: """Cleanup all Gluetun containers on exit.""" # Get instances without holding the lock during cleanup with cleanup_lock: @@ -37,7 +37,7 @@ def cleanup_all_gluetun_containers(): pass -def register_cleanup(): +def register_cleanup() -> None: """Add the cleanup handlers to atexit, one time only.""" global cleanup_registered with cleanup_lock: @@ -168,8 +168,8 @@ class Gluetun(Proxy): auth_user: Optional[str] = None, auth_password: Optional[str] = None, verify_ip: bool = True, - **kwargs, - ): + **kwargs: Any, + ) -> None: """ Initialise the Gluetun proxy provider that can use more than one VPN provider. @@ -202,7 +202,7 @@ class Gluetun(Proxy): self.verify_ip = verify_ip # Track active containers: {query_key: {"container_name": ..., "port": ..., ...}} - self.active_containers = {} + self.active_containers: dict[str, dict[str, Any]] = {} self._port_lock = threading.Lock() @@ -367,7 +367,7 @@ class Gluetun(Proxy): del self.active_containers[query_key] raise RuntimeError(f"Failed to create Gluetun container: {e}") - def cleanup(self): + def cleanup(self) -> None: """Stop and remove all managed Gluetun containers.""" container_count = len(self.active_containers) @@ -438,7 +438,7 @@ class Gluetun(Proxy): "org": container.get("ip_org"), } - def validate_provider_config(self, provider_name: str, config: dict): + def validate_provider_config(self, provider_name: str, config: dict) -> None: """Validate a provider's configuration.""" vpn_type = config.get("vpn_type", "wireguard").lower() credentials = config.get("credentials", {}) @@ -600,7 +600,7 @@ class Gluetun(Proxy): country: Optional[str] = None, city: Optional[str] = None, hostname: Optional[str] = None, - ): + ) -> None: """ Make and start a Gluetun Docker container. @@ -1073,7 +1073,7 @@ class Gluetun(Proxy): except (subprocess.TimeoutExpired, FileNotFoundError): return "" - def verify_container(self, query_key: str, max_retries: int = 3): + def verify_container(self, query_key: str, max_retries: int = 3) -> None: """ Make sure that the container's VPN IP and region are correct with an ipinfo.io lookup. @@ -1271,7 +1271,7 @@ class Gluetun(Proxy): except (subprocess.TimeoutExpired, FileNotFoundError): return False - def remove_container(self, container_name: str): + def remove_container(self, container_name: str) -> None: """Stop and remove a Docker container.""" try: if self.auto_cleanup: @@ -1311,7 +1311,7 @@ class Gluetun(Proxy): return f"http://{self.auth_user}:{self.auth_password}@localhost:{port}" return f"http://localhost:{port}" - def __del__(self): + def __del__(self) -> None: """Cleanup containers on object destruction.""" if hasattr(self, "auto_cleanup") and self.auto_cleanup: try: diff --git a/unshackle/core/proxies/hola.py b/unshackle/core/proxies/hola.py index 26e99f7f..c21ee5d9 100644 --- a/unshackle/core/proxies/hola.py +++ b/unshackle/core/proxies/hola.py @@ -8,13 +8,13 @@ from unshackle.core.proxies.proxy import Proxy class Hola(Proxy): - def __init__(self): + def __init__(self) -> None: """ Proxy provider that uses Hola's direct connections through the hola-proxy binary. """ - self.binary = binaries.HolaProxy - if not self.binary: + if not binaries.HolaProxy: raise EnvironmentError("hola-proxy executable not found but is required for the Hola proxy provider.") + self.binary = binaries.HolaProxy self.countries = self.get_countries() diff --git a/unshackle/core/proxies/nordvpn.py b/unshackle/core/proxies/nordvpn.py index 9a4691a1..095f49dd 100644 --- a/unshackle/core/proxies/nordvpn.py +++ b/unshackle/core/proxies/nordvpn.py @@ -11,7 +11,7 @@ from unshackle.core.utilities import COUNTRY_CODE_ALIASES class NordVPN(Proxy): - def __init__(self, username: str, password: str, server_map: Optional[dict[str, Union[int, str]]] = None): + def __init__(self, username: str, password: str, server_map: Optional[dict[str, Union[int, str]]] = None) -> None: """ Proxy provider that uses NordVPN Service Credentials. @@ -139,6 +139,7 @@ class NordVPN(Proxy): [by_id is None or country["id"] == int(by_id), by_code is None or country["code"] == by_code.upper()] ): return country + return None @staticmethod def get_all_servers(country_id: int) -> list[dict]: diff --git a/unshackle/core/proxies/proton.py b/unshackle/core/proxies/proton.py index 967c42d8..a9bc6552 100644 --- a/unshackle/core/proxies/proton.py +++ b/unshackle/core/proxies/proton.py @@ -9,6 +9,7 @@ import sys import time from pathlib import Path from typing import Optional +from urllib.parse import quote import requests @@ -62,7 +63,7 @@ class ProtonVPN(Proxy): cache_path: Optional[str] = None, timeout: float = 10.0, enable: bool = False, - ): + ) -> None: self.server_map = {str(k).lower().strip(): str(v).lower().strip() for k, v in (server_map or {}).items()} self.timeout = timeout self.enable = enable @@ -109,7 +110,7 @@ class ProtonVPN(Proxy): self.last_host = host = str(server.get("Domain") or "") port = self.SECURE_CORE_PORT if int(server.get("Features") or 0) & self.SECURE_CORE_FEATURE else self.PROXY_PORT # Username (a JWT) and Password are URL-safe, but quote so requests rebuilds Basic auth cleanly - user, pwd = (requests.utils.quote(c, safe="") for c in creds) + user, pwd = (quote(c, safe="") for c in creds) return f"https://{user}:{pwd}@{host}:{port}" def last_connection_display(self) -> Optional[str]: diff --git a/unshackle/core/proxies/proxy.py b/unshackle/core/proxies/proxy.py index bb974552..fd00fa95 100644 --- a/unshackle/core/proxies/proxy.py +++ b/unshackle/core/proxies/proxy.py @@ -1,10 +1,10 @@ from abc import abstractmethod -from typing import Optional +from typing import Any, Optional class Proxy: @abstractmethod - def __init__(self, **kwargs): + def __init__(self, **kwargs: Any) -> None: """ The constructor initializes the Service using passed configuration data. diff --git a/unshackle/core/proxies/surfsharkvpn.py b/unshackle/core/proxies/surfsharkvpn.py index 0eaaffb6..cd00401b 100644 --- a/unshackle/core/proxies/surfsharkvpn.py +++ b/unshackle/core/proxies/surfsharkvpn.py @@ -83,7 +83,7 @@ def pinned(ip: str, hostname: str) -> int: class SurfsharkVPN(Proxy): - def __init__(self, username: str, password: str, server_map: Optional[dict[str, str]] = None): + def __init__(self, username: str, password: str, server_map: Optional[dict[str, str]] = None) -> None: """ Proxy provider that uses SurfsharkVPN Service Credentials. diff --git a/unshackle/core/proxies/windscribevpn.py b/unshackle/core/proxies/windscribevpn.py index 7603d7ae..16dd7b65 100644 --- a/unshackle/core/proxies/windscribevpn.py +++ b/unshackle/core/proxies/windscribevpn.py @@ -10,7 +10,7 @@ from unshackle.core.proxies.proxy import Proxy class WindscribeVPN(Proxy): - def __init__(self, username: str, password: str, server_map: Optional[dict[str, str]] = None): + def __init__(self, username: str, password: str, server_map: Optional[dict[str, str]] = None) -> None: """ Proxy provider that uses WindscribeVPN Service Credentials. diff --git a/unshackle/core/remote_service.py b/unshackle/core/remote_service.py index 7c317864..2572206a 100644 --- a/unshackle/core/remote_service.py +++ b/unshackle/core/remote_service.py @@ -615,12 +615,14 @@ def match_track(remote_track: Track, local_tracks: list) -> Optional[Track]: if str(lt.id) == remote_id: return lt + if not isinstance(remote_track, (Video, Audio, Subtitle)): + return None for lt in local_tracks: if type(lt).__name__ != type(remote_track).__name__: continue if lt.codec != remote_track.codec or str(lt.language) != str(remote_track.language): continue - if hasattr(lt, "width") and hasattr(remote_track, "width"): + if isinstance(remote_track, Video): if ( lt.width == remote_track.width and lt.height == remote_track.height @@ -628,10 +630,10 @@ def match_track(remote_track: Track, local_tracks: list) -> Optional[Track]: and same_bitrate(lt.bitrate, remote_track.bitrate) ): return lt - elif hasattr(lt, "channels") and hasattr(remote_track, "channels"): + elif isinstance(remote_track, Audio): if same_bitrate(lt.bitrate, remote_track.bitrate): return lt - elif hasattr(lt, "forced"): + elif isinstance(remote_track, Subtitle): if lt.forced == remote_track.forced and lt.sdh == remote_track.sdh: return lt return None @@ -1074,7 +1076,7 @@ class ServerVault(Vault): report_now = True - def __init__(self, service: "RemoteService"): + def __init__(self, service: "RemoteService") -> None: super().__init__("server_cdm", no_push=True) self.service = service @@ -1167,7 +1169,7 @@ class RemoteService: total=5, backoff_factor=0.2, status_forcelist=[429, 500, 502, 503, 504], - retry_after_max=MAX_BACKOFF, + retry_after_max=int(MAX_BACKOFF), ), pool_maxsize=64, pool_block=True, @@ -1326,7 +1328,7 @@ class RemoteService: f"No local proxy for the server's account region ({', '.join(regions)}), " f"you are in {client_region}" ) - elif not no_proxy and proxy: + elif not no_proxy and proxy and self.ctx.parent: resolved_proxy = resolve_remote_proxy_arg(server_proxy_arg(self.ctx.parent.params), providers) if resolved_proxy: create_data["proxy"] = resolved_proxy @@ -1804,7 +1806,7 @@ class RemoteService: self._titles = Movies(titles_list) return self._titles - def get_titles_cached(self, title_id: str = None) -> Titles_T: + def get_titles_cached(self, title_id: Optional[str] = None) -> Titles_T: """Apply the client's local title_map to titles fetched from the remote server. Lets users rename titles for remote services they do not have installed locally. @@ -2084,9 +2086,8 @@ class RemoteService: from unshackle.core.drm import Widevine - kid_uuids = [UUID(hex=k) for k in kid_hexes] WIDEVINE_SYSTEM_ID = UUID("edef8ba9-79d6-4ace-a3c8-27dcd51d21ed") - dummy_pssh = WvPSSH.new(system_id=WIDEVINE_SYSTEM_ID, key_ids=kid_uuids) + dummy_pssh = WvPSSH.new(system_id=WIDEVINE_SYSTEM_ID, key_ids=[UUID(hex=k) for k in kid_hexes]) # no kid: the stub names several KIDs, and the first is not the track's own return Widevine(pssh=dummy_pssh) diff --git a/unshackle/core/search_result.py b/unshackle/core/search_result.py index 537b5d50..1ffe6ca7 100644 --- a/unshackle/core/search_result.py +++ b/unshackle/core/search_result.py @@ -9,7 +9,7 @@ class SearchResult: description: Optional[str] = None, label: Optional[str] = None, url: Optional[str] = None, - ): + ) -> None: """ A search result for any title type a service can give. diff --git a/unshackle/core/service.py b/unshackle/core/service.py index 506dc56f..19a3da46 100644 --- a/unshackle/core/service.py +++ b/unshackle/core/service.py @@ -5,7 +5,7 @@ from collections.abc import Callable, Generator from dataclasses import dataclass, field from http.cookiejar import CookieJar from pathlib import Path -from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union +from typing import TYPE_CHECKING, Any, Optional, Sequence, TypeVar, Union if TYPE_CHECKING: from unshackle.core.api.input_bridge import InputBridge @@ -84,10 +84,10 @@ class TimeoutHTTPAdapter(HTTPAdapter): self.default_timeout = timeout super().__init__(*args, **kwargs) - def send(self, request: Any, **kwargs: Any) -> Any: + def send(self, request: requests.PreparedRequest, *args: Any, **kwargs: Any) -> requests.Response: if kwargs.get("timeout") is None: kwargs["timeout"] = getattr(self, "default_timeout", DEFAULT_TIMEOUT) - return super().send(request, **kwargs) + return super().send(request, *args, **kwargs) def grow_session_pool(session: Any, size: int) -> None: @@ -174,7 +174,7 @@ class Service(metaclass=ABCMeta): # one of them set keeps the server CDM even when the remote client picked its own device. SERVER_DEVICE_OPTIONS: tuple[str, ...] = () - def __init__(self, ctx: click.Context): + def __init__(self, ctx: click.Context) -> None: console.print(Padding(Rule(f"[rule.text]Service: {self.__class__.__name__}"), (1, 2))) self.config = ctx.obj.config @@ -187,7 +187,7 @@ class Service(metaclass=ABCMeta): self.cache_dir = config.directories.cache / self.__class__.__name__ self.ctx = ctx - self.credential = None # Will be set in authenticate() + self.credential: Optional[Credential] = None # Will be set in authenticate() self.current_region = None # Will be set based on proxy/geolocation self._input_bridge: Optional[InputBridge] = None @@ -349,7 +349,7 @@ class Service(metaclass=ABCMeta): held = {Tracks.signature(track) for track in all_tracks} all_tracks.add([track for track in tracks if Tracks.signature(track) not in held], warn_only=True) - codecs = self.track_request.codecs or [None] + codecs: Sequence[Optional[Video.Codec]] = self.track_request.codecs or [None] ranges = self.track_request.ranges or [Video.Range.SDR] for range_val in ranges: @@ -419,7 +419,7 @@ class Service(metaclass=ABCMeta): status_forcelist=STATUS_FORCELIST, allowed_methods=RETRY_METHODS, respect_retry_after_header=True, - retry_after_max=MAX_BACKOFF, + retry_after_max=int(MAX_BACKOFF), ), pool_connections=POOL_MAX_SIZE, pool_maxsize=POOL_MAX_SIZE, @@ -510,7 +510,7 @@ class Service(metaclass=ABCMeta): def get_widevine_service_certificate( self, *, challenge: bytes, title: Title_T, track: AnyTrack - ) -> Union[bytes, str]: + ) -> Optional[Union[bytes, str]]: """ Get the Widevine Service Certificate used for Privacy Mode. @@ -604,7 +604,7 @@ class Service(metaclass=ABCMeta): This can be useful to store information on each title that you need later, like any sub-asset IDs, or such. """ - def get_titles_cached(self, title_id: str = None) -> Titles_T: + def get_titles_cached(self, title_id: Optional[str] = None) -> Titles_T: """ Cached wrapper around get_titles() to reduce redundant API calls. diff --git a/unshackle/core/service_repo.py b/unshackle/core/service_repo.py index a9a3f6e6..e8a8cd9d 100644 --- a/unshackle/core/service_repo.py +++ b/unshackle/core/service_repo.py @@ -16,7 +16,7 @@ log = logging.getLogger("services") class DirtyServiceRepo(Exception): """A clone has uncommitted or unpushed local changes. unshackle refuses to refresh it, to protect them.""" - def __init__(self, path: Path): + def __init__(self, path: Path) -> None: self.path = path super().__init__(str(path)) diff --git a/unshackle/core/services.py b/unshackle/core/services.py index 5deb1127..aca57408 100644 --- a/unshackle/core/services.py +++ b/unshackle/core/services.py @@ -88,7 +88,7 @@ def compiled_import_hint(missing: str, path: Path) -> str: ) -def load_service(path: Path) -> object: +def load_service(path: Path) -> type[Service]: """Load one Service module, returning its tag-named class. Raises a concise, single-line error naming the Service and the real cause so @@ -108,7 +108,7 @@ def load_service(path: Path) -> object: ) from e -def load_services(paths: list[Path]) -> tuple[dict[str, object], list[str]]: +def load_services(paths: list[Path]) -> tuple[dict[str, type[Service]], list[str]]: """Load every Service, returning the good ones plus a list of load errors. Importing this module must never raise: several commands import it, and Python @@ -116,7 +116,7 @@ def load_services(paths: list[Path]) -> tuple[dict[str, object], list[str]]: every command. Instead we collect failures and let the caller surface them once, cleanly, when it uses the services. """ - modules: dict[str, object] = {} + modules: dict[str, type[Service]] = {} errors: list[str] = [] for path in paths: try: @@ -126,7 +126,7 @@ def load_services(paths: list[Path]) -> tuple[dict[str, object], list[str]]: return modules, errors -def register_service_binaries(modules: dict[str, object]) -> None: +def register_service_binaries(modules: dict[str, type[Service]]) -> None: """Register custom binary dependencies declared by services via get_binaries().""" for tag, service_cls in modules.items(): @@ -570,7 +570,9 @@ class Services(click.Group): def import_cli(ctx: click.Context, title: str, **kwargs: object) -> object: from unshackle.core.import_service import ImportService - import_file = ctx.params.get("import_file") or (ctx.parent and ctx.parent.params.get("import_file")) + import_file = ctx.params.get("import_file") or ( + ctx.parent.params.get("import_file") if ctx.parent else None + ) return ImportService(ctx, tag, title, import_file) return import_cli @@ -612,7 +614,7 @@ class Services(click.Group): return value @staticmethod - def load(tag: str) -> Service: + def load(tag: str) -> type[Service]: """Load a Service module by Service tag.""" module = MODULES.get(tag) if module: diff --git a/unshackle/core/session.py b/unshackle/core/session.py index a5ca50ce..dedebb92 100644 --- a/unshackle/core/session.py +++ b/unshackle/core/session.py @@ -13,11 +13,11 @@ from collections.abc import Iterator, MutableMapping from datetime import datetime, timezone from email.utils import parsedate_to_datetime from http.cookiejar import CookieJar -from typing import Any, Optional +from typing import Any, Optional, Union from urllib.parse import urlencode, urlparse, urlunparse import rnet -from requests import HTTPError, Request, Response +from requests import HTTPError, Request, Response, Session from requests.cookies import create_cookie from requests.structures import CaseInsensitiveDict @@ -102,13 +102,13 @@ class RnetResponseHeaders(MutableMapping): def __len__(self) -> int: return self._map.keys_len() - def get(self, key: str, default: Optional[str] = None) -> Optional[str]: + def get(self, key: str, default: Optional[str] = None) -> Optional[str]: # type: ignore[override] val = self._map.get(key) if val is None: return default return self.decode(val) - def items(self) -> list[tuple[str, str]]: + def items(self) -> list[tuple[str, str]]: # type: ignore[override] return [(self.decode(k), self.decode(v)) for k, v in self._map.items()] @@ -196,7 +196,7 @@ class RnetResponse: if not self.ok: raise HTTPError( f"{self.status_code} {self.reason}: {self.url}", - response=self, + response=self, # type: ignore[arg-type] ) def iter_content(self, chunk_size: Optional[int] = None) -> Iterator[bytes]: @@ -358,7 +358,7 @@ class RnetCookieAdapter(MutableMapping): jar.set_cookie(cookie) return jar - def update(self, other: Any = None, **kwargs: Any) -> None: + def update(self, other: Any = None, /, **kwargs: Any) -> None: if other is None: other = {} if isinstance(other, CookieJar): @@ -385,7 +385,7 @@ class RnetCookieAdapter(MutableMapping): self._flat[name] = value self.set_cookie_on_client("https://localhost", name, value) - def get( + def get( # type: ignore[override] self, name: str, default: Optional[str] = None, domain: Optional[str] = None, path: Optional[str] = None ) -> Optional[str]: if domain and domain in self._cookies: @@ -484,13 +484,13 @@ class RnetCookieAdapter(MutableMapping): if self._client is not None: self._client.clear_cookies() - def items(self) -> list[tuple[str, str]]: + def items(self) -> list[tuple[str, str]]: # type: ignore[override] return list(self._flat.items()) - def keys(self) -> list[str]: + def keys(self) -> list[str]: # type: ignore[override] return list(self._flat.keys()) - def values(self) -> list[str]: + def values(self) -> list[str]: # type: ignore[override] return list(self._flat.values()) @@ -892,7 +892,10 @@ class RnetSession: if response.status_code not in self.status_forcelist: return response - last_exception = HTTPError(f"Received status code: {response.status_code}", response=response) + last_exception = HTTPError( + f"Received status code: {response.status_code}", + response=response, # type: ignore[arg-type] + ) failure = failure_reason(last_exception) except self.catch_exceptions as e: @@ -953,8 +956,8 @@ class RnetSession: send_kwargs: dict[str, Any] = {} if req.headers: send_kwargs["headers"] = dict(req.headers) - if req.body: - send_kwargs["data"] = req.body + if req.data: + send_kwargs["data"] = req.data if req.json: send_kwargs["json"] = req.json @@ -970,6 +973,9 @@ class RnetSession: pass +AnySession = Union[Session, RnetSession] + + def session( browser: Optional[str] = None, **kwargs: Any, diff --git a/unshackle/core/title_cacher.py b/unshackle/core/title_cacher.py index 0262e637..afd126fb 100644 --- a/unshackle/core/title_cacher.py +++ b/unshackle/core/title_cacher.py @@ -3,10 +3,11 @@ from __future__ import annotations import hashlib import logging from datetime import datetime, timedelta -from typing import Optional +from typing import Callable, Optional from unshackle.core.cacher import Cacher from unshackle.core.config import config +from unshackle.core.credential import Credential from unshackle.core.titles import Titles_T @@ -21,7 +22,7 @@ class TitleCacher: - Cache hit/miss statistics for debugging """ - def __init__(self, service_name: str): + def __init__(self, service_name: str) -> None: self.service_name = service_name self.log = logging.getLogger(f"{service_name}.TitleCache") self.cacher = Cacher(service_name) @@ -61,12 +62,12 @@ class TitleCacher: def get_cached_titles( self, title_id: str, - fetch_function, + fetch_function: Callable[[], Titles_T], region: Optional[str] = None, account_hash: Optional[str] = None, no_cache: bool = False, reset_cache: bool = False, - ) -> Optional[Titles_T]: + ) -> Titles_T: """ Get titles from the cache, or fetch them from the API with a fallback. @@ -115,7 +116,7 @@ class TitleCacher: return titles except Exception as e: - if cache and cache.data: + if cache and cache.data and cache.expiration: current_time = datetime.now() max_retention_time = cache.expiration + timedelta( seconds=config.title_cache_max_retention - config.title_cache_time @@ -140,7 +141,7 @@ class TitleCacher: raise - def clear_all_title_cache(self): + def clear_all_title_cache(self) -> None: """Clear all title caches for this service.""" cache_dir = config.directories.cache / self.service_name if cache_dir.exists(): @@ -265,7 +266,7 @@ class TitleCacher: self.log.debug(f"Cached {provider_name} data for {title_id}") -def get_account_hash(credential) -> Optional[str]: +def get_account_hash(credential: Optional[Credential]) -> Optional[str]: """ Make a hash for account identification. diff --git a/unshackle/core/titles/episode.py b/unshackle/core/titles/episode.py index d810f1db..edd1b484 100644 --- a/unshackle/core/titles/episode.py +++ b/unshackle/core/titles/episode.py @@ -213,7 +213,7 @@ class Episode(Title): def get_filename(self, media_info: MediaInfo, folder: bool = False, show_service: bool = True) -> str: if folder: - template = config.get_folder_template("series") + template: Optional[str] = config.get_folder_template("series") if template: # all parts of an episode land in the same season folder context = self.build_template_context(media_info, show_service, include_part=False) @@ -268,7 +268,7 @@ class Episode(Title): class Series(SortedKeyList, ABC): - def __init__(self, iterable: Optional[Iterable] = None): + def __init__(self, iterable: Optional[Iterable] = None) -> None: # part slots before year so existing ties still resolve by year super().__init__(iterable, key=lambda x: (x.season, x.number, x.part or 0, x.year or 0)) diff --git a/unshackle/core/titles/movie.py b/unshackle/core/titles/movie.py index fdb718f7..df581fd0 100644 --- a/unshackle/core/titles/movie.py +++ b/unshackle/core/titles/movie.py @@ -61,7 +61,7 @@ class Movie(Title): def get_filename(self, media_info: MediaInfo, folder: bool = False, show_service: bool = True) -> str: if folder: - template = config.get_folder_template("movies") + template: Optional[str] = config.get_folder_template("movies") if template: context = self.build_template_context(media_info, show_service) spacer = detect_spacer(template) # one style for the whole path @@ -86,7 +86,7 @@ class Movie(Title): class Movies(SortedKeyList, ABC): - def __init__(self, iterable: Optional[Iterable] = None): + def __init__(self, iterable: Optional[Iterable] = None) -> None: super().__init__(iterable, key=lambda x: x.year or 0) def __str__(self) -> str: diff --git a/unshackle/core/titles/music.py b/unshackle/core/titles/music.py index 65702ef7..5e096ece 100644 --- a/unshackle/core/titles/music.py +++ b/unshackle/core/titles/music.py @@ -230,7 +230,7 @@ class Song(Title): class Music(SortedKeyList, ABC): """A grouped music release, such as an album, EP, single, compilation, or playlist.""" - def __new__(cls, *args: Any, **kwargs: Any): + def __new__(cls, *args: Any, **kwargs: Any) -> "Music": return super().__new__(cls) def __init__( @@ -246,7 +246,7 @@ class Music(SortedKeyList, ABC): total_duration: Optional[int] = None, owner: Optional[str] = None, description: Optional[str] = None, - ): + ) -> None: if not isinstance(kind, str): raise TypeError(f"Expected kind to be a str, not {kind!r}") if title is not None and not isinstance(title, str): diff --git a/unshackle/core/titles/title.py b/unshackle/core/titles/title.py index ebfbe0d4..341efc08 100644 --- a/unshackle/core/titles/title.py +++ b/unshackle/core/titles/title.py @@ -58,13 +58,15 @@ class Title: self.id = id_ self.service = service self.language = language - self.data = data + self.data: Any = data self.anime: Optional[bool] = None self.daily: Optional[bool] = None self.tracks = Tracks() - def __eq__(self, other: Title) -> bool: + def __eq__(self, other: object) -> bool: + if not isinstance(other, Title): + return NotImplemented return self.id == other.id def build_base_template_context(self, media_info: MediaInfo, show_service: bool = True) -> dict: diff --git a/unshackle/core/tracks/attachment.py b/unshackle/core/tracks/attachment.py index 619199b3..091fce1c 100644 --- a/unshackle/core/tracks/attachment.py +++ b/unshackle/core/tracks/attachment.py @@ -8,15 +8,11 @@ from typing import Optional, Union from urllib.parse import urlparse from zlib import crc32 -import requests - from unshackle.core.config import config from unshackle.core.constants import DOWNLOAD_LICENCE_ONLY -from unshackle.core.session import RnetSession +from unshackle.core.session import AnySession from unshackle.core.utils import run_proxy -AnySession = Union[requests.Session, RnetSession] - class Attachment: def __init__( @@ -27,7 +23,7 @@ class Attachment: mime_type: Optional[str] = None, description: Optional[str] = None, session: Optional[AnySession] = None, - ): + ) -> None: """ Make a new Attachment. @@ -137,7 +133,7 @@ class Attachment: ) def __str__(self) -> str: - return " | ".join(filter(bool, ["ATT", self.name, self.mime_type, self.description])) + return " | ".join(filter(None, ["ATT", self.name, self.mime_type, self.description])) def to_dict(self) -> dict[str, Optional[str]]: """Serialise a URL-backed attachment for export/import.""" diff --git a/unshackle/core/tracks/audio.py b/unshackle/core/tracks/audio.py index 78348569..9a8f9e2f 100644 --- a/unshackle/core/tracks/audio.py +++ b/unshackle/core/tracks/audio.py @@ -111,7 +111,7 @@ class Audio(Track): joc: Optional[int] = None, descriptive: Union[bool, int] = False, **kwargs: Any, - ): + ) -> None: """ Make a new Audio track object. @@ -196,14 +196,14 @@ class Audio(Track): def __str__(self) -> str: return " | ".join( filter( - bool, + None, [ "AUD", f"[{self.codec.value}]" if self.codec else None, str(self.language), ", ".join( filter( - bool, + None, [ str(self.channels) if self.channels else None, "Atmos" if self.atmos else None, diff --git a/unshackle/core/tracks/chapter.py b/unshackle/core/tracks/chapter.py index 20e21e86..ef24ebf9 100644 --- a/unshackle/core/tracks/chapter.py +++ b/unshackle/core/tracks/chapter.py @@ -8,7 +8,7 @@ TIMESTAMP_FORMAT = re.compile(r"^(?P\d{2}):(?P\d{2}):(?P\d class Chapter: - def __init__(self, timestamp: Union[str, int, float], name: Optional[str] = None): + def __init__(self, timestamp: Union[str, int, float], name: Optional[str] = None) -> None: """ Make a new Chapter with a Timestamp and optional name. @@ -31,6 +31,9 @@ class Chapter: raise TypeError(f"Expected name to be {str}, not {type(name)}") if not isinstance(timestamp, str): + hours: float + minutes: float + remainder: float if isinstance(timestamp, int): # ms hours, remainder = divmod(timestamp, 1000 * 60 * 60) minutes, remainder = divmod(remainder, 1000 * 60) @@ -47,8 +50,8 @@ class Chapter: if not timestamp_m: raise ValueError(f"The timestamp format is invalid: {timestamp}") - hour, minute, second, ms = timestamp_m.groups() - if not ms: + hour, minute, second, ms_group = timestamp_m.groups() + if not ms_group: timestamp += ".000" self.timestamp = timestamp @@ -60,7 +63,7 @@ class Chapter: ) def __str__(self) -> str: - return " | ".join(filter(bool, ["CHP", self.timestamp, self.name])) + return " | ".join(filter(None, ["CHP", self.timestamp, self.name])) @property def id(self) -> str: diff --git a/unshackle/core/tracks/chapters.py b/unshackle/core/tracks/chapters.py index 71898bd4..78a42f8a 100644 --- a/unshackle/core/tracks/chapters.py +++ b/unshackle/core/tracks/chapters.py @@ -17,7 +17,7 @@ OGM_SIMPLE_LINE_2_FORMAT = re.compile(r"^CHAPTER(?P\d+)NAME=(?P.*) class Chapters(SortedKeyList, ABC): - def __init__(self, iterable: Optional[Iterable[Chapter]] = None): + def __init__(self, iterable: Optional[Iterable[Chapter]] = None) -> None: super().__init__(key=lambda x: x.timestamp or 0) for chapter in iterable or []: self.add(chapter) diff --git a/unshackle/core/tracks/dv_fixup.py b/unshackle/core/tracks/dv_fixup.py index b52389e4..75e1bc49 100644 --- a/unshackle/core/tracks/dv_fixup.py +++ b/unshackle/core/tracks/dv_fixup.py @@ -45,6 +45,7 @@ class DVFixup: """Demux the elementary HEVC so mkvmerge exposes its DV RPU. Returns the elementary-stream path, or the original source on any failure so muxing can proceed with the as-downloaded file.""" + assert FFMPEG and self.video.path, "checked in __init__" source = Path(self.video.path) height = self.video.height or 0 console.print(Padding(Rule(f"[rule.text]DV Composite Fixup ({height}p)"), (1, 2))) diff --git a/unshackle/core/tracks/hybrid.py b/unshackle/core/tracks/hybrid.py index f5a703c8..73533828 100644 --- a/unshackle/core/tracks/hybrid.py +++ b/unshackle/core/tracks/hybrid.py @@ -6,7 +6,7 @@ import subprocess from collections import Counter from concurrent import futures from pathlib import Path -from typing import Optional +from typing import TYPE_CHECKING, Any, Optional from rich.padding import Padding from rich.rule import Rule @@ -18,6 +18,9 @@ from unshackle.core.utilities import log_event from unshackle.core.utils import dovi from unshackle.core.utils.subprocess import run_step +if TYPE_CHECKING: + from unshackle.core.tracks.video import Video + # HDR10+ histogram slots that form the maxRGB distribution, as (DistributionIndex, fraction of pixels). # Index 99 holds the 99.98% percentile. Index 5 (a luminance-Y value) and index 10 (the share of pixels # at or under 100 nits) are not part of the distribution. @@ -111,7 +114,7 @@ def hdr10plus_shots(metadata: dict, active_fraction: float = 1.0) -> list[dict]: class Hybrid: - def __init__(self, videos, source, output_name: Optional[str] = None) -> None: + def __init__(self, videos: list["Video"], source: str, output_name: Optional[str] = None) -> None: self.log = logging.getLogger("hybrid") """ @@ -237,7 +240,7 @@ class Hybrid: Path.unlink(config.directories.temp / "hdr10.mkv") Path.unlink(config.directories.temp / "dv.mkv") - def extract_stream(self, save_path, type_): + def extract_stream(self, save_path: Path, type_: str) -> None: output = Path(config.directories.temp / f"{type_}.hevc") try: run_step( @@ -263,7 +266,7 @@ class Hybrid: success=True, ) - def extract_rpu(self, video, untouched=False): + def extract_rpu(self, video: "Video", untouched: bool = False) -> None: if os.path.isfile(config.directories.temp / "RPU.bin") or os.path.isfile( config.directories.temp / "RPU_UNT.bin" ): @@ -512,7 +515,7 @@ class Hybrid: max_fall = max_cll return max_mdl, min_mdl, max_cll, max_fall - def level_6(self): + def level_6(self) -> None: """Edit RPU Level 6 values using the static L6 luminance data from the RPU.""" if os.path.isfile(config.directories.temp / "RPU_L6.bin"): self.rpu_file = "RPU_L6.bin" @@ -531,7 +534,7 @@ class Hybrid: raise ValueError("Failed reading RPU metadata for Level 6 values") blocks = rpu_info.get("vdr_dm_data", {}).get("cmv29_metadata", {}).get("ext_metadata_blocks", []) - level6 = next((b["Level6"] for b in blocks if isinstance(b, dict) and "Level6" in b), {}) + level6: dict[str, Any] = next((b["Level6"] for b in blocks if isinstance(b, dict) and "Level6" in b), {}) max_cll = level6.get("max_content_light_level") max_fall = level6.get("max_frame_average_light_level") @@ -623,7 +626,7 @@ class Hybrid: ) self.rpu_file = "RPU_L6.bin" - def injecting(self): + def injecting(self) -> None: if os.path.isfile(config.directories.temp / self.hevc_file): return @@ -656,7 +659,7 @@ class Hybrid: success=True, ) - def extract_hdr10plus(self, video): + def extract_hdr10plus(self, video: "Video") -> None: """Extract HDR10+ metadata from the video track""" if os.path.isfile(config.directories.temp / self.hdr10plus_file): return @@ -706,7 +709,7 @@ class Hybrid: success=True, ) - def probe_hdr_metadata(self): + def probe_hdr_metadata(self) -> tuple[int, int, int, int]: """Extract mastering display and Content Light Level metadata from the HDR10 track with FFprobe. Returns (max_mdl, min_mdl, max_cll, max_fall) in dovi_tool level6 units: @@ -793,9 +796,8 @@ class Hybrid: # one entry per frame: a feature film is hundreds of MB of objects, freed before dovi_tool runs del metadata - max_mdl, min_mdl, max_cll, max_fall = self.probe_hdr_metadata() - max_mdl, min_mdl, max_cll, max_fall = self.sanitize_l6(max_mdl, min_mdl, max_cll, max_fall) - generator_config = { + max_mdl, min_mdl, max_cll, max_fall = self.sanitize_l6(*self.probe_hdr_metadata()) + generator_config: dict[str, Any] = { "cm_version": "V29", "level6": { "max_display_mastering_luminance": max_mdl, diff --git a/unshackle/core/tracks/subtitle.py b/unshackle/core/tracks/subtitle.py index 9cfde2e5..0f86ed8c 100644 --- a/unshackle/core/tracks/subtitle.py +++ b/unshackle/core/tracks/subtitle.py @@ -13,7 +13,6 @@ from typing import Any, Callable, Iterable, Optional, Sequence, Union import click import pycaption -import requests from construct import Container from langcodes import Language from pycaption import Caption, CaptionList, CaptionNode, WebVTTReader @@ -24,6 +23,7 @@ from subtitle_filter import Subtitles from unshackle.core import binaries from unshackle.core.config import config +from unshackle.core.session import AnySession from unshackle.core.tracks.track import Track from unshackle.core.utilities import PYMP4_LOCK, try_ensure_utf8, with_script from unshackle.core.utils.subprocess import log_tool_run @@ -112,7 +112,7 @@ class Subtitle(Track): sdh: bool = False, forced: bool = False, **kwargs: Any, - ): + ) -> None: """ Make a new Subtitle track object. @@ -230,7 +230,7 @@ class Subtitle(Track): def __str__(self) -> str: return " | ".join( filter( - bool, + None, [ "SUB", f"[{self.codec.value}]" if self.codec else None, @@ -253,7 +253,7 @@ class Subtitle(Track): def download( self, - session: requests.Session, + session: AnySession, prepare_drm: partial, max_workers: Optional[int] = None, progress: Optional[partial] = None, @@ -263,7 +263,7 @@ class Subtitle(Track): proxy_download: Optional[str] = None, adaptive_workers: bool = False, download_processes: int = 1, - ): + ) -> None: super().download( session, prepare_drm, @@ -522,7 +522,7 @@ class Subtitle(Track): return text lines = text.split("\n") - cues = [] + cues: list[dict[str, Any]] = [] header_lines = [] in_header = True i = 0 @@ -746,7 +746,7 @@ class Subtitle(Track): text = try_ensure_utf8(data).decode("utf8") caption_set = pycaption.SRTReader().read(text) elif codec == Subtitle.Codec.fTTML: - caption_lists: dict[str, pycaption.CaptionList] = defaultdict(pycaption.CaptionList) + caption_lists: dict[Optional[str], pycaption.CaptionList] = defaultdict(pycaption.CaptionList) with PYMP4_LOCK: boxes = MP4.parse_stream(BytesIO(data)) for segment in ( @@ -754,7 +754,7 @@ class Subtitle(Track): ): for lang in segment.get_languages(): caption_lists[lang].extend(segment.get_captions(lang)) - caption_set: pycaption.CaptionSet = pycaption.CaptionSet(caption_lists) + caption_set = pycaption.CaptionSet(caption_lists) elif codec == Subtitle.Codec.TimedTextMarkupLang: text = try_ensure_utf8(data).decode("utf8") text = text.replace("tt:", "") @@ -767,10 +767,10 @@ class Subtitle(Track): text = re.sub(r"-(\d+(?:\.\d+)?)(px|em|%|c|pt)", r"0\2", text) caption_set = pycaption.DFXPReader().read(text) elif codec == Subtitle.Codec.fVTT: - caption_lists: dict[str, pycaption.CaptionList] = defaultdict(pycaption.CaptionList) + caption_lists = defaultdict(pycaption.CaptionList) caption_list, language = Subtitle.merge_segmented_wvtt(data) caption_lists[language] = caption_list - caption_set: pycaption.CaptionSet = pycaption.CaptionSet(caption_lists) + caption_set = pycaption.CaptionSet(caption_lists) elif codec == Subtitle.Codec.WebVTT: text = try_ensure_utf8(data).decode("utf8") text = Subtitle.sanitize_broken_webvtt(text) @@ -837,7 +837,7 @@ class Subtitle(Track): return "\n".join(sanitized_lines) @staticmethod - def space_webvtt_headers(data: Union[str, bytes]): + def space_webvtt_headers(data: Union[str, bytes]) -> str: """ Space out the WEBVTT Headers from Captions. @@ -857,7 +857,7 @@ class Subtitle(Track): return text @staticmethod - def merge_same_cues(caption_set: pycaption.CaptionSet): + def merge_same_cues(caption_set: pycaption.CaptionSet) -> None: """Merge captions with the same timecodes and text as one in-place.""" for lang in caption_set.get_languages(): captions = caption_set.get_captions(lang) @@ -882,7 +882,7 @@ class Subtitle(Track): caption_set.set_captions(lang, merged_captions) @staticmethod - def filter_unwanted_cues(caption_set: pycaption.CaptionSet): + def filter_unwanted_cues(caption_set: pycaption.CaptionSet) -> None: """ Filter out subtitle cues containing only   or whitespace. """ @@ -924,7 +924,7 @@ class Subtitle(Track): base_time = 0 # > trun saw_trun_box = False - samples = [] + samples: list[Any] = [] def flatten_boxes(box: Container) -> Iterable[Container]: for child in box: @@ -1117,6 +1117,7 @@ class Subtitle(Track): if binaries.SubtitleEdit and use_subtitleedit: from unshackle.core.tracks.subtitle_convert import SUBTITLE_EDIT_FORMATS, subtitleedit_args + assert self.codec is not None output_format = SUBTITLE_EDIT_FORMATS.get(self.codec, self.codec.name.lower()) se_start = time.monotonic() subprocess.run( @@ -1177,6 +1178,7 @@ class Subtitle(Track): from unshackle.core.tracks.subtitle_convert import SUBTITLE_EDIT_FORMATS, subtitleedit_args + assert self.codec is not None output_format = SUBTITLE_EDIT_FORMATS.get(self.codec, self.codec.name.lower()) se_start = time.monotonic() diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index 3fd6536c..da4290d6 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -26,11 +26,11 @@ from unshackle.core import binaries from unshackle.core.cdm.detect import is_playready_cdm, is_widevine_cdm from unshackle.core.config import config from unshackle.core.constants import DOWNLOAD_ALL_DRM, DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, DownloadCancelled -from unshackle.core.downloaders import requests +from unshackle.core.downloaders.requests import requests from unshackle.core.drm import DRM_T, BlowfishStripe, ClearKeyCENC, PlayReady, Widevine from unshackle.core.drm.verify import decrypt_track, has_encrypted_entry, senc_protects_samples from unshackle.core.events import events -from unshackle.core.session import MAX_BACKOFF, RnetSession, utf8_text +from unshackle.core.session import MAX_BACKOFF, AnySession, RnetSession, utf8_text from unshackle.core.temp import restore_segments, stash_segments from unshackle.core.tracks import resume from unshackle.core.utilities import get_boxes, language_name, log_event, try_ensure_utf8 @@ -45,7 +45,7 @@ DRM_PREFERENCE_TYPES: dict[str, Union[type[Widevine], type[PlayReady]]] = { } -def direct_session(session: Union[Session, "RnetSession"], proxy: Optional[str] = None) -> Session: +def direct_session(session: AnySession, proxy: Optional[str] = None) -> Session: """requests.Session with copied headers/cookies and no proxy, or only ``proxy`` when given.""" new = Session() if proxy: @@ -82,7 +82,7 @@ def direct_session(session: Union[Session, "RnetSession"], proxy: Optional[str] total=5, backoff_factor=0.2, status_forcelist=[429, 500, 502, 503, 504], - retry_after_max=MAX_BACKOFF, + retry_after_max=int(MAX_BACKOFF), ), pool_maxsize=64, pool_block=True, @@ -390,7 +390,7 @@ class DownloadContext: save_path: Path save_dir: Path progress: partial - session: Optional[Union[Session, "RnetSession"]] = None + session: Optional[AnySession] = None proxy: Optional[str] = None max_workers: Optional[int] = None adaptive_workers: bool = False @@ -398,7 +398,7 @@ class DownloadContext: license_widevine: Optional[Callable] = None cdm: Optional[object] = None - def ensure_session(self) -> Union[Session, "RnetSession"]: + def ensure_session(self) -> AnySession: """Return the HTTP session, or a new ``Session`` if none was set.""" session = self.session if not session: @@ -428,7 +428,7 @@ class Track: drm: Optional[Iterable[DRM_T]] = None, edition: Optional[str] = None, note: Optional[str] = None, - session: Optional[Union[Session, "RnetSession"]] = None, + session: Optional[AnySession] = None, downloader: Optional[Callable] = None, downloader_args: Optional[dict] = None, from_file: Optional[Path] = None, @@ -454,7 +454,7 @@ class Track: raise TypeError(f"Expected edition to be a {str}, {list}, or None, not {type(edition)}") if not isinstance(note, (str, type(None))): raise TypeError(f"Expected note to be a {str}, not {type(note)}") - if not isinstance(downloader, (Callable, type(None))): + if downloader is not None and not callable(downloader): raise TypeError(f"Expected downloader to be a {Callable}, not {type(downloader)}") if not isinstance(downloader_args, (dict, type(None))): raise TypeError(f"Expected downloader_args to be a {dict}, not {type(downloader_args)}") @@ -463,7 +463,7 @@ class Track: if not isinstance(data, (dict, defaultdict, type(None))): raise TypeError(f"Expected data to be a {dict} or {defaultdict}, not {type(data)}") - invalid_urls = ", ".join(set(type(x) for x in url if not isinstance(x, str))) + invalid_urls = ", ".join({str(type(x)) for x in url if not isinstance(x, str)}) if invalid_urls: raise TypeError(f"Expected all items in url to be a {str}, but found {invalid_urls}") @@ -497,7 +497,8 @@ class Track: if not id_: this = copy(self) - this.url = self.url.rsplit("?", maxsplit=1)[0] + if isinstance(self.url, str): + this.url = self.url.rsplit("?", maxsplit=1)[0] checksum = crc32(repr(this).encode("utf8")) id_ = hex(checksum)[2:] @@ -527,7 +528,7 @@ class Track: self.sample_start = (starts[first] - starts[0]) / timescale return slice(first, max(bisect_left(starts, high) + self.sample_tail, first + 1)) - def sample_byte_range(self, session: Union[Session, "RnetSession"]) -> Optional[tuple[bytes, str]]: + def sample_byte_range(self, session: AnySession) -> Optional[tuple[bytes, str]]: """Find the bytes of a one-file track that hold the sample window. Reads the start of the file for a ``sidx`` box. Returns the bytes before that box, which hold @@ -535,12 +536,13 @@ class Track: """ if self.sample_bytes or not isinstance(self.url, str): return self.sample_bytes + url = self.url head = b"" def read_to(end: int) -> None: nonlocal head if end > len(head): - res = session.get(self.url, headers={"Range": f"bytes={len(head)}-{end - 1}"}) + res = session.get(url, headers={"Range": f"bytes={len(head)}-{end - 1}"}) if res.ok: head += res.content @@ -565,7 +567,7 @@ class Track: return self.sample_bytes return None - def can_sample(self, session: Union[Session, "RnetSession"]) -> bool: + def can_sample(self, session: AnySession) -> bool: """Return False when a sample cannot cut this track, because it is one file with no segment index.""" if self.descriptor == self.Descriptor.URL: return bool(self.sample_byte_range(session)) @@ -682,7 +684,7 @@ class Track: def download( self, - session: Union[Session, "RnetSession"], + session: AnySession, prepare_drm: partial, max_workers: Optional[int] = None, progress: Optional[partial] = None, @@ -692,14 +694,18 @@ class Track: proxy_download: Optional[str] = None, adaptive_workers: bool = False, download_processes: int = 1, - ): + ) -> None: """Download and optionally Decrypt this Track. For a URL-descriptor Video or Audio track with no `drm` set, unshackle probes the DRM from the track's init data and stores it on the track, so a service need not declare `drm` itself. """ from unshackle.core.manifests import DASH, HLS, ISM + from unshackle.core.tracks.audio import Audio + from unshackle.core.tracks.subtitle import Subtitle + from unshackle.core.tracks.video import Video + assert isinstance(self, (Video, Audio, Subtitle)) and progress is not None if DOWNLOAD_LICENCE_ONLY.is_set(): progress(downloaded="[yellow]SKIPPING") @@ -722,7 +728,8 @@ class Track: track_type = self.__class__.__name__ save_path = config.directories.temp / f"{track_type}_{self.id}.mp4" - if track_type == "Subtitle": + if isinstance(self, Subtitle): + assert self.codec save_path = save_path.with_suffix(f".{self.codec.extension}") if self.descriptor != self.Descriptor.URL: @@ -769,7 +776,7 @@ class Track: cleanup() try: - manifest_parsers = { + manifest_parsers: dict[Track.Descriptor, Union[type[DASH], type[HLS], type[ISM]]] = { self.Descriptor.HLS: HLS, self.Descriptor.DASH: DASH, self.Descriptor.ISM: ISM, @@ -893,7 +900,7 @@ class Track: if has_encrypted_sample_entry(save_path): self.needs_repack = True - if track_type == "Subtitle" and self.codec.name not in ("fVTT", "fTTML"): + if isinstance(self, Subtitle) and self.codec and self.codec.name not in ("fVTT", "fTTML"): track_data = self.path.read_bytes() track_data = try_ensure_utf8(track_data) track_data = ( @@ -928,6 +935,7 @@ class Track: return if not DOWNLOAD_LICENCE_ONLY.is_set(): + assert self.path is not None if self.path.stat().st_size <= 3: # Empty UTF-8 BOM == 3 bytes raise IOError("Download failed, the downloaded file is empty.") @@ -1016,7 +1024,7 @@ class Track: from unshackle.core.tracks.video import Video track_type = data.get("type") - builders = {"Video": Video, "Audio": Audio, "Subtitle": Subtitle} + builders: dict[Optional[str], type[Track]] = {"Video": Video, "Audio": Audio, "Subtitle": Subtitle} builder = builders.get(track_type) if builder is None: raise ValueError(f"Cannot reconstruct unsupported track type: {track_type!r}") @@ -1075,9 +1083,9 @@ class Track: playready[0].absorb(*playready[1:]) return playready[0] - return self.drm[0] + return next(iter(self.drm)) - def get_key_id(self, init_data: Optional[bytes] = None, *args, **kwargs) -> Optional[UUID]: + def get_key_id(self, init_data: Optional[bytes] = None, *args: Any, **kwargs: Any) -> Optional[UUID]: """ Probe the DRM encryption Key ID (KID) for this specific track. @@ -1117,8 +1125,9 @@ class Track: tenc = uuid_box.data if tenc.key_ID.int != 0: return tenc.key_ID + return None - def load_drm_if_needed(self, service=None) -> bool: + def load_drm_if_needed(self, service: Optional[Any] = None) -> bool: """ Load DRM information for this track if the parser deferred it. @@ -1164,6 +1173,7 @@ class Track: from pywidevine.pssh import PSSH as WV_PSSH session = getattr(self, "session", None) or run_proxy.fallback_session() + assert isinstance(self.url, str) response = session.get(self.url) playlist = m3u8.loads(utf8_text(response), self.url) @@ -1177,7 +1187,7 @@ class Track: fmt = key.keyformat.lower() if fmt == WidevineCdm.urn: pssh_b64 = key.uri.split(",")[-1] - drm = Widevine(pssh=WV_PSSH(pssh_b64)) + drm: DRM_T = Widevine(pssh=WV_PSSH(pssh_b64)) drm_list.append(drm) elif fmt in {f"urn:uuid:{PR_PSSH.SYSTEM_ID}", "com.microsoft.playready"}: pssh_b64 = key.uri.split(",")[-1] @@ -1199,7 +1209,7 @@ class Track: maximum_size: int = 20000, url: Optional[str] = None, byte_range: Optional[str] = None, - session: Optional[Session] = None, + session: Optional[AnySession] = None, ) -> bytes: """ Get the Track's initial segment data. @@ -1235,7 +1245,7 @@ class Track: if not url: if self.descriptor != self.Descriptor.URL: raise ValueError(f"An explicit URL must be provided for {self.descriptor.name} tracks") - if not self.url: + if not self.url or not isinstance(self.url, str): raise ValueError("An explicit URL must be provided as the track has no URL") url = self.url @@ -1262,6 +1272,7 @@ class Track: if range_test.status_code == 206: byte_range = f"0-{content_length - 1}" + init_data: Optional[bytes] if byte_range: res = session.get(url=url, headers={"Range": f"bytes={byte_range}"}) res.raise_for_status() @@ -1294,15 +1305,16 @@ class Track: if not self.path or not self.path.exists(): raise ValueError("Cannot repackage a Track that has not been downloaded.") - if not binaries.FFMPEG: + ffmpeg_bin = binaries.FFMPEG + if not ffmpeg_bin: 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}_repack") - def ffmpeg(extra_args: list[str] = None, bsf: Optional[str] = None): - args = [ - binaries.FFMPEG, + def ffmpeg(extra_args: Optional[list[str]] = None, bsf: Optional[str] = None) -> None: + args: list[Union[str, Path]] = [ + ffmpeg_bin, "-nostdin", "-hide_banner", "-loglevel", diff --git a/unshackle/core/tracks/tracks.py b/unshackle/core/tracks/tracks.py index 12a3f006..697dd2cc 100644 --- a/unshackle/core/tracks/tracks.py +++ b/unshackle/core/tracks/tracks.py @@ -77,7 +77,7 @@ class Tracks: Tracks, Sequence[Union[AnyTrack, Chapter, Chapters, Attachment]], Track, Chapter, Chapters, Attachment ], manifest_url: Optional[str] = None, - ): + ) -> None: self.videos: list[Video] = [] self.audio: list[Audio] = [] self.subtitles: list[Subtitle] = [] @@ -109,7 +109,7 @@ class Tracks: ) def __str__(self) -> str: - rep = {Video: [], Audio: [], Subtitle: [], Chapter: [], Attachment: []} + rep: dict[type, list[str]] = {Video: [], Audio: [], Subtitle: [], Chapter: [], Attachment: []} tracks = [*list(self), *self.chapters] for track in sorted(tracks, key=lambda t: self.TRACK_ORDER_MAP[type(t)]): @@ -125,14 +125,11 @@ class Tracks: ) rep[type(track)].append(str(track)) - for type_ in list(rep): - if not rep[type_]: - del rep[type_] - continue - rep[type_] = "\n".join([rep[type_][0]] + [f"├─ {x}" for x in rep[type_][1:-1]] + [f"└─ {rep[type_][-1]}"]) - rep = "\n".join(list(rep.values())) - - return rep + return "\n".join( + "\n".join([lines[0]] + [f"├─ {x}" for x in lines[1:-1]] + [f"└─ {lines[-1]}"]) + for lines in rep.values() + if lines + ) def tree(self, add_progress: bool = False) -> tuple[Tree, list[Callable[..., None]]]: all_tracks = [*list(self), *self.chapters, *self.attachments] @@ -236,7 +233,12 @@ class Tracks: def add( self, tracks: Union[ - Tracks, Sequence[Union[AnyTrack, Chapter, Chapters, Attachment]], Track, Chapter, Chapters, Attachment + Tracks, + Sequence[Union[Tracks, Sequence[Any], Track, Chapter, Chapters, Attachment]], + Track, + Chapter, + Chapters, + Attachment, ], warn_only: bool = False, ) -> None: @@ -485,7 +487,9 @@ class Tracks: if video not in non_hybrid_selected and not video.native_hybrid: video.hybrid_base_only = True - def select_hybrid(self, tracks, quality, worst: bool = False, generate_dv: bool = False): + def select_hybrid( + self, tracks: list[Video], quality: list[int], worst: bool = False, generate_dv: bool = False + ) -> Callable[[Video], bool]: """A filter for the hybrid ingredients of each resolution. It sets ``native_hybrid`` on every track. A track that already carries DV and an HDR10 or HDR10+ base (DV+HDR10P, HDR10+DV, HDR10P+DV) is a finished @@ -526,12 +530,14 @@ class Tracks: for res in quality: candidates = [v for v in base_tracks if at_resolution(v, res)] if candidates: - chosen = pick(candidates, key=lambda v: v.bitrate) + chosen = pick(candidates, key=lambda v: v.bitrate or 0) base_selected.append(chosen) - lowest_dv = min(dv_tracks, key=lambda v: v.height) if dv_tracks and (base_selected or not natives) else None + lowest_dv = ( + min(dv_tracks, key=lambda v: v.height or 0) if dv_tracks and (base_selected or not natives) else None + ) - def select(x): + def select(x: Video) -> bool: if x in base_selected or x.native_hybrid: return True if lowest_dv and x is lowest_dv: @@ -1017,13 +1023,13 @@ class Tracks: if output_path is None: output_path = ( self.videos[0].path.with_suffix(".muxed.mkv") - if self.videos + if self.videos and self.videos[0].path else self.audio[0].path.with_suffix(".muxed.mka") - if self.audio + if self.audio and self.audio[0].path else self.subtitles[0].path.with_suffix(".muxed.mks") - if self.subtitles + if self.subtitles and self.subtitles[0].path else chapters_path.with_suffix(".muxed.mkv") - if self.chapters + if self.chapters and chapters_path else None ) if not output_path: @@ -1063,12 +1069,13 @@ class Tracks: warnings = [] mux_start_time = time.monotonic() p = subprocess.Popen(full_command, text=True, encoding="utf-8", errors="replace", stdout=subprocess.PIPE) + assert p.stdout is not None for line in iter(p.stdout.readline, ""): if line.startswith("#GUI#error") or line.startswith("#GUI#warning"): errors.append(line) if line.startswith("#GUI#warning"): warnings.append(line.strip()) - if "progress" in line: + if progress and "progress" in line: progress(total=100, completed=int(line.strip()[14:-1])) returncode = p.wait() diff --git a/unshackle/core/tracks/video.py b/unshackle/core/tracks/video.py index a6765584..fc129b58 100644 --- a/unshackle/core/tracks/video.py +++ b/unshackle/core/tracks/video.py @@ -185,7 +185,7 @@ class Video(Track): # The codebase is currently agnostic to either, so a manual conversion to 6 is done. transfer = 6 - def safe(enum_cls, value): + def safe(enum_cls: type[Enum], value: int) -> Any: try: return enum_cls(value) except ValueError: @@ -391,18 +391,18 @@ class Video(Track): def __str__(self) -> str: return " | ".join( filter( - bool, + None, [ "VID", - "[" + (", ".join(filter(bool, [self.codec.value if self.codec else None, self.range_label]))) + "]", + "[" + (", ".join(filter(None, [self.codec.value if self.codec else None, self.range_label]))) + "]", str(self.language), ", ".join( filter( - bool, + None, [ " @ ".join( filter( - bool, + None, [ f"{self.width}x{self.height}" if self.width and self.height else None, f"{self.bitrate // 1000} kb/s" if self.bitrate else None, @@ -667,12 +667,13 @@ class Video(Track): raise EnvironmentError("ccextractor executable was not found.") out_path = Path(out_path) + video_path, ccextractor = self.path, binaries.CCExtractor def run_ccextractor(fallback: str) -> bool: cc_start = time.monotonic() try: subprocess.run( - [binaries.CCExtractor, "-trim", "-nobom", "-noru", "-ru1", "-o", out_path, self.path], + [ccextractor, "-trim", "-nobom", "-noru", "-ru1", "-o", out_path, video_path], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, @@ -687,7 +688,7 @@ class Video(Track): ) if e.returncode < 0: logging.getLogger("Video").warning( - f"ccextractor crashed (signal {-e.returncode}) on {self.path.name}; {fallback}" + f"ccextractor crashed (signal {-e.returncode}) on {video_path.name}; {fallback}" ) return out_path.exists() if e.returncode != 10: # 10 = No captions found diff --git a/unshackle/core/utilities.py b/unshackle/core/utilities.py index c0565a58..1e4685ad 100644 --- a/unshackle/core/utilities.py +++ b/unshackle/core/utilities.py @@ -19,7 +19,7 @@ from datetime import datetime, timezone from functools import cache from pathlib import Path from types import ModuleType -from typing import Any, Callable, Mapping, Optional, Sequence, Union +from typing import Any, Callable, Iterator, Mapping, Optional, Sequence, Union from urllib.parse import ParseResult, urlparse from uuid import uuid4 @@ -101,12 +101,12 @@ def import_module_by_path(path: Path) -> ModuleType: # compute package hierarchy for relative import support if path.is_relative_to(config.directories.core_dir): - name = [] + parts: list[str] = [] _path = path.parent while _path.stem != config.directories.core_dir.stem: - name.append(_path.stem) + parts.append(_path.stem) _path = _path.parent - name = ".".join([config.directories.core_dir.stem] + name[::-1]) + name = ".".join([config.directories.core_dir.stem] + parts[::-1]) else: # is outside the src package if str(path.parent.parent) not in sys.path: @@ -114,6 +114,7 @@ def import_module_by_path(path: Path) -> ModuleType: name = path.parent.stem spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) sys.modules[name] = module spec.loader.exec_module(module) @@ -544,7 +545,7 @@ def get_boxes(data: bytes, box_type: bytes, as_bytes: bool = False) -> Box: # t raise e -def ap_case(text: str, keep_spaces: bool = False, stop_words: tuple[str] = None) -> str: +def ap_case(text: str, keep_spaces: bool = False, stop_words: Optional[tuple[str, ...]] = None) -> str: """ Convert a string to title case using AP/APA style. Based on https://github.com/words/ap-style-title-case @@ -758,6 +759,7 @@ def get_extension(value: Union[str, Path, ParseResult]) -> Optional[str]: ext = os.path.splitext(value_parsed.path)[1] if ext and ext != ".": return ext + return None def extract_font_family(font_path: Path) -> Optional[str]: @@ -810,6 +812,8 @@ def get_windows_fonts() -> dict[str, Path]: Returns: Dictionary mapping font family names to their file paths """ + if sys.platform != "win32": + raise OSError("Windows font registry is only available on Windows") import winreg with winreg.ConnectRegistry(None, winreg.HKEY_LOCAL_MACHINE) as reg: @@ -850,7 +854,7 @@ def get_unix_fonts() -> dict[str, Path]: Dictionary mapping font family names to their file paths """ log = logging.getLogger("get_system_fonts") - fonts = {} + fonts: dict[str, Path] = {} font_dirs = [ Path("/usr/share/fonts"), @@ -965,7 +969,7 @@ def find_font_with_fallbacks(font_name: str, system_fonts: dict[str, Path]) -> O # Mapping of font families to system packages that provide them -FONT_PACKAGES = { +FONT_PACKAGES: dict[str, dict[str, Any]] = { "liberation": { "debian": "fonts-liberation fonts-liberation2", "fonts": ["Liberation Sans", "Liberation Serif", "Liberation Mono"], @@ -995,7 +999,7 @@ def suggest_font_packages(missing_fonts: list[str]) -> dict[str, list[str]]: Returns: Dictionary mapping package names to lists of fonts they would give """ - suggestions = {} + suggestions: dict[str, list[str]] = {} # Check which fonts from aliases would help needed_aliases = set() @@ -1048,7 +1052,7 @@ class DebugLogger: This makes it easy to parse, filter, and analyse logs programmatically. """ - def __init__(self, log_path: Optional[Path] = None, enabled: bool = False, log_keys: bool = False): + def __init__(self, log_path: Optional[Path] = None, enabled: bool = False, log_keys: bool = False) -> None: """ Initialise the debug logger. @@ -1063,12 +1067,12 @@ class DebugLogger: self.file_handle = None self.log_keys = log_keys - if self.enabled: + if self.enabled and self.log_path is not None: self.log_path.parent.mkdir(parents=True, exist_ok=True) self.file_handle = open(self.log_path, "a", encoding="utf-8") self.log_session_start() - def log_session_start(self): + def log_session_start(self) -> None: """Log the ``session_start`` entry, with environment information.""" import platform @@ -1095,13 +1099,13 @@ class DebugLogger: message: str = "", context: Optional[dict[str, Any]] = None, service: Optional[str] = None, - error: Optional[Exception] = None, + error: Optional[BaseException] = None, request: Optional[dict[str, Any]] = None, response: Optional[dict[str, Any]] = None, duration_ms: Optional[float] = None, success: Optional[bool] = None, - **kwargs, - ): + **kwargs: Any, + ) -> None: """ Log a structured JSON entry. @@ -1121,7 +1125,7 @@ class DebugLogger: if not self.enabled or not self.file_handle: return - entry = { + entry: dict[str, Any] = { "timestamp": datetime.now(timezone.utc).isoformat(), "session_id": self.session_id, "level": level, @@ -1219,7 +1223,7 @@ class DebugLogger: except Exception: return f"[{type(data).__name__}]" - def log_operation_start(self, operation: str, **kwargs) -> str: + def log_operation_start(self, operation: str, **kwargs: Any) -> str: """ Log the start of an operation and return an operation ID. @@ -1241,8 +1245,13 @@ class DebugLogger: return op_id def log_operation_end( - self, operation: str, operation_id: str, success: bool = True, duration_ms: Optional[float] = None, **kwargs - ): + self, + operation: str, + operation_id: str, + success: bool = True, + duration_ms: Optional[float] = None, + **kwargs: Any, + ) -> None: """ Log the end of an operation. @@ -1263,7 +1272,7 @@ class DebugLogger: **kwargs, ) - def log_service_call(self, method: str, url: str, **kwargs): + def log_service_call(self, method: str, url: str, **kwargs: Any) -> None: """ Log a service API call. @@ -1278,7 +1287,7 @@ class DebugLogger: level=level, operation="service_call", message=message, request={"method": method, "url": url, **kwargs} ) - def log_drm_operation(self, drm_type: str, operation: str, **kwargs): + def log_drm_operation(self, drm_type: str, operation: str, **kwargs: Any) -> None: """ Log a DRM operation (PSSH extraction, license request, content key retrieval). @@ -1291,7 +1300,7 @@ class DebugLogger: level = kwargs.pop("level", "DEBUG") self.log(level=level, operation=f"drm_{operation}", message=message, drm_type=drm_type, **kwargs) - def log_vault_query(self, vault_name: str, operation: str, **kwargs): + def log_vault_query(self, vault_name: str, operation: str, **kwargs: Any) -> None: """ Log a vault query operation. @@ -1310,7 +1319,7 @@ class DebugLogger: **kwargs, ) - def log_error(self, operation: str, error: Exception, **kwargs): + def log_error(self, operation: str, error: Exception, **kwargs: Any) -> None: """ Log an error with full context. @@ -1328,7 +1337,7 @@ class DebugLogger: **kwargs, ) - def close(self): + def close(self) -> None: """Close the log file and clean up resources.""" if self.file_handle: self.log(level="INFO", operation="session_end", message="Debug logging session ended") @@ -1358,7 +1367,7 @@ def log_event(operation: str, *, level: str = "DEBUG", message: str = "", **kwar @contextlib.contextmanager -def timed_operation(operation: str, *, level: str = "DEBUG", message: str = "", **kwargs: Any): +def timed_operation(operation: str, *, level: str = "DEBUG", message: str = "", **kwargs: Any) -> Iterator[None]: """Time a block and log ``operation`` once it finishes, with ``duration_ms``. Logs at ``level`` with ``success=True`` on normal exit, or at ERROR with the exception and @@ -1398,7 +1407,7 @@ def timed_operation(operation: str, *, level: str = "DEBUG", message: str = "", ) -def init_debug_logger(log_path: Optional[Path] = None, enabled: bool = False, log_keys: bool = False): +def init_debug_logger(log_path: Optional[Path] = None, enabled: bool = False, log_keys: bool = False) -> None: """ Initialise the global debug logger. @@ -1413,7 +1422,7 @@ def init_debug_logger(log_path: Optional[Path] = None, enabled: bool = False, lo debug_logger = DebugLogger(log_path=log_path, enabled=enabled, log_keys=log_keys) -def close_debug_logger(): +def close_debug_logger() -> None: """Close the global debug logger.""" global debug_logger if debug_logger: diff --git a/unshackle/core/utils/bitrate.py b/unshackle/core/utils/bitrate.py index efdea62e..5fbbad0a 100644 --- a/unshackle/core/utils/bitrate.py +++ b/unshackle/core/utils/bitrate.py @@ -7,12 +7,13 @@ import time from collections import OrderedDict, defaultdict from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable, Hashable, Optional, Union +from typing import TYPE_CHECKING, Callable, Hashable, Optional, Sequence, Union from urllib.parse import urljoin from requests import Session from unshackle.core.binaries import FFProbe +from unshackle.core.constants import TrackT from unshackle.core.session import RnetSession, utf8_text from unshackle.core.utils.subprocess import log_tool_run @@ -115,11 +116,11 @@ def measure_real_bitrate( def apply_real_bitrates( - tracks: list["Track"], + tracks: list[TrackT], session: Union[Session, RnetSession], *, log: logging.Logger, - group_key: Callable[["Track"], Hashable], + group_key: Callable[[TrackT], Hashable], per_group: int = 5, workers: int = 8, ) -> None: @@ -134,7 +135,7 @@ def apply_real_bitrates( outrank a probed one), it probes the next track too, until the probed set is safely above the rest. Unprobed tracks keep their manifest-declared bitrate. """ - groups: defaultdict[Hashable, list["Track"]] = defaultdict(list) + groups: defaultdict[Hashable, list[TrackT]] = defaultdict(list) for track in tracks: groups[group_key(track)].append(track) for group in groups.values(): @@ -155,7 +156,7 @@ def apply_real_bitrates( def probe_batch( - tracks: list["Track"], + tracks: Sequence["Track"], session: Union[Session, RnetSession], *, log: logging.Logger, diff --git a/unshackle/core/utils/click_types.py b/unshackle/core/utils/click_types.py index b76aae8d..5e4797f8 100644 --- a/unshackle/core/utils/click_types.py +++ b/unshackle/core/utils/click_types.py @@ -1,10 +1,9 @@ import re from datetime import date, timedelta -from typing import Any, Optional, Union +from typing import Any, Optional, Sequence, Union import click from click.shell_completion import CompletionItem -from pywidevine.cdm import Cdm as WidevineCdm from unshackle.core.constants import DEFAULT_SAMPLE from unshackle.core.tracks.audio import Audio @@ -88,7 +87,7 @@ class SubtitleCodecChoice(click.Choice): - Common aliases: srt (for SubRip) """ - def __init__(self, codec_enum): + def __init__(self, codec_enum: type[Subtitle.Codec]) -> None: self.codec_enum = codec_enum choices = [] aliases = {} @@ -135,7 +134,9 @@ class SubtitleCodecChoice(click.Choice): self.aliases = aliases super().__init__(choices, case_sensitive=False) - def convert(self, value: Any, param: Optional[click.Parameter] = None, ctx: Optional[click.Context] = None): + def convert( + self, value: Any, param: Optional[click.Parameter] = None, ctx: Optional[click.Context] = None + ) -> Union[Subtitle.Codec, str, None]: if not value: return None @@ -160,11 +161,11 @@ class ContextData: def __init__( self, config: dict, - cdm: WidevineCdm, + cdm: Any, proxy_providers: list, profile: Optional[str] = None, cdm_entry: bool = False, - ): + ) -> None: self.config = config self.cdm = cdm self.proxy_providers = proxy_providers @@ -342,7 +343,7 @@ class QualityList(click.ParamType): def convert( self, - value: Union[str, int, list[Union[str, int]]], + value: Union[str, int, Sequence[Union[str, int]]], param: Optional[click.Parameter] = None, ctx: Optional[click.Context] = None, ) -> list[int]: diff --git a/unshackle/core/utils/dovi.py b/unshackle/core/utils/dovi.py index 9b16f346..6f7918a6 100644 --- a/unshackle/core/utils/dovi.py +++ b/unshackle/core/utils/dovi.py @@ -16,7 +16,7 @@ import json import subprocess import time from pathlib import Path -from typing import Optional +from typing import Optional, Union from unshackle.core import binaries from unshackle.core.utils.subprocess import log_tool_run, run_step @@ -137,7 +137,7 @@ def generate( ) -> bytes: """Build a DV RPU from a dovi_tool generator config, and from hdr10plus_tool JSON when given.""" tool = require_dovi_tool() - extra = ["--hdr10plus-json", hdr10plus_json] if hdr10plus_json else [] + extra: list[Union[str, Path]] = ["--hdr10plus-json", hdr10plus_json] if hdr10plus_json else [] return run_step( [tool, "generate", "-j", config_json, *extra, "-o", output], status=status, diff --git a/unshackle/core/utils/gen_esn.py b/unshackle/core/utils/gen_esn.py index e58c4a32..933d7aca 100644 --- a/unshackle/core/utils/gen_esn.py +++ b/unshackle/core/utils/gen_esn.py @@ -6,11 +6,11 @@ from datetime import datetime, timedelta log = logging.getLogger("NF-ESN") -def chrome_esn_generator(): +def chrome_esn_generator() -> str: ESN_GEN = "".join(random.choice("0123456789ABCDEF") for _ in range(30)) esn_file = ".esn" - def gen_file(): + def gen_file() -> None: with open(esn_file, "w") as file: file.write(f"NFCDIE-03-{ESN_GEN}") diff --git a/unshackle/core/utils/redact.py b/unshackle/core/utils/redact.py index bc1a47cd..69af12e6 100644 --- a/unshackle/core/utils/redact.py +++ b/unshackle/core/utils/redact.py @@ -6,7 +6,7 @@ import logging import re import sys from pathlib import Path -from typing import Iterable, Optional +from typing import Iterable, Optional, overload from urllib.parse import urlsplit, urlunsplit REDACTED = "***" @@ -116,6 +116,10 @@ def path_redaction_enabled() -> bool: return True +@overload +def redact_path(text: str) -> str: ... +@overload +def redact_path(text: None) -> None: ... def redact_path(text: Optional[str]) -> Optional[str]: """Replace local base-directory prefixes (install root, venv, home) in ``text`` with tokens. diff --git a/unshackle/core/utils/selector.py b/unshackle/core/utils/selector.py index fe73c082..88b40b09 100644 --- a/unshackle/core/utils/selector.py +++ b/unshackle/core/utils/selector.py @@ -1,4 +1,5 @@ import sys +from typing import Any, Literal, overload import click from rich.console import Group @@ -10,8 +11,6 @@ from rich.text import Text from unshackle.core.console import console IS_WINDOWS = sys.platform == "win32" -if IS_WINDOWS: - import msvcrt class Selector: @@ -27,9 +26,9 @@ class Selector: text_style: str = "text", page_size: int = 8, minimal_count: int = 0, - dependencies: dict[int, list[int]] = None, + dependencies: dict[int, list[int]] | None = None, collapse_on_start: bool = False, - ): + ) -> None: """ Initialise the Selector. @@ -49,16 +48,16 @@ class Selector: self.minimal_count = minimal_count self.dependencies = dependencies or {} - self.child_to_parent = {} + self.child_to_parent: dict[int, int] = {} for parent, children in self.dependencies.items(): for child in children: self.child_to_parent[child] = parent self.cursor_index = 0 - self.selected_indices = set() + self.selected_indices: set[int] = set() self.scroll_offset = 0 - self.expanded_headers = set() + self.expanded_headers: set[int] = set() if not collapse_on_start: self.expanded_headers.update(self.dependencies.keys()) @@ -77,7 +76,7 @@ class Selector: visible.append(idx) return visible - def get_renderable(self): + def get_renderable(self) -> Padding: """ Constructs and returns the renderable object (Table + Info) for the current state. """ @@ -144,7 +143,7 @@ class Selector: return Padding(Group(table, info_text), (0, 5)) - def move_cursor(self, delta: int): + def move_cursor(self, delta: int) -> None: """ Moves the cursor up or down through VISIBLE items only. """ @@ -160,7 +159,7 @@ class Selector: new_visual_idx = (current_visual_idx + delta) % len(visible_indices) self.cursor_index = visible_indices[new_visual_idx] - def change_page(self, delta: int): + def change_page(self, delta: int) -> None: """ Changes the current page view by the specified delta (previous/next page). """ @@ -190,7 +189,7 @@ class Selector: except IndexError: pass - def toggle_selection(self): + def toggle_selection(self) -> None: """ Toggles the selection state of the item currently under the cursor. """ @@ -206,7 +205,7 @@ class Selector: else: self.selected_indices.difference_update(target_indices) - def toggle_expand(self): + def toggle_expand(self) -> None: """ Expands or collapses the current header. """ @@ -216,7 +215,7 @@ class Selector: else: self.expanded_headers.add(self.cursor_index) - def toggle_expand_all(self): + def toggle_expand_all(self) -> None: """ Toggles expansion state of ALL headers. If every header shows its children -> Collapse all. @@ -230,7 +229,7 @@ class Selector: else: self.expanded_headers = all_headers.copy() - def toggle_all(self): + def toggle_all(self) -> None: """ Toggles the selection of all items. """ @@ -239,11 +238,15 @@ class Selector: else: self.selected_indices = set(range(len(self.options))) - def get_input_windows(self): + def get_input_windows(self) -> str | None: """ Captures and parses keyboard input on Windows systems using msvcrt. Returns command strings such as 'UP', 'DOWN', or 'ENTER'. """ + if sys.platform != "win32": + raise OSError("msvcrt key input is only available on Windows") + import msvcrt + key = msvcrt.getch() # Ctrl+C (0x03) or ESC (0x1b) if key == b"\x03" or key == b"\x1b": @@ -290,7 +293,7 @@ class Selector: return "RIGHT" return None - def get_input_unix(self): + def get_input_unix(self) -> str | None: """ Captures and parses keyboard input on Unix/Linux systems using click.getchar(). Returns command strings such as 'UP', 'DOWN', or 'ENTER'. @@ -390,6 +393,31 @@ class Selector: return [] +@overload +def select_multiple( + options: list[str], + minimal_count: int = ..., + page_size: int = ..., + return_indices: Literal[True] = ..., + cursor_style: str = ..., + collapse_on_start: bool = ..., + **kwargs: Any, +) -> list[int]: ... + + +@overload +def select_multiple( + options: list[str], + minimal_count: int = ..., + page_size: int = ..., + *, + return_indices: Literal[False], + cursor_style: str = ..., + collapse_on_start: bool = ..., + **kwargs: Any, +) -> list[str]: ... + + def select_multiple( options: list[str], minimal_count: int = 1, @@ -397,8 +425,8 @@ def select_multiple( return_indices: bool = True, cursor_style: str = "pink", collapse_on_start: bool = False, - **kwargs, -) -> list[int]: + **kwargs: Any, +) -> list[int] | list[str]: """ Drop-in replacement using custom Selector with global console. diff --git a/unshackle/core/utils/sslciphers.py b/unshackle/core/utils/sslciphers.py index d85aa320..583da32d 100644 --- a/unshackle/core/utils/sslciphers.py +++ b/unshackle/core/utils/sslciphers.py @@ -1,5 +1,5 @@ import ssl -from typing import Optional +from typing import Any, Optional from requests.adapters import HTTPAdapter @@ -46,7 +46,7 @@ class SSLCiphers(HTTPAdapter): __attrs__ = [*HTTPAdapter.__attrs__, "_ssl_context"] - def __init__(self, cipher_list: Optional[str] = None, security_level: int = 0, *args, **kwargs): + def __init__(self, cipher_list: Optional[str] = None, security_level: int = 0, *args: Any, **kwargs: Any) -> None: if cipher_list: if not isinstance(cipher_list, str): raise TypeError(f"Expected cipher_list to be a str, not {cipher_list!r}") @@ -70,10 +70,10 @@ class SSLCiphers(HTTPAdapter): self._ssl_context = ctx super().__init__(*args, **kwargs) - def init_poolmanager(self, *args, **kwargs): + def init_poolmanager(self, *args: Any, **kwargs: Any) -> None: kwargs["ssl_context"] = self._ssl_context return super().init_poolmanager(*args, **kwargs) - def proxy_manager_for(self, *args, **kwargs): + def proxy_manager_for(self, *args: Any, **kwargs: Any) -> Any: kwargs["ssl_context"] = self._ssl_context return super().proxy_manager_for(*args, **kwargs) diff --git a/unshackle/core/utils/subprocess.py b/unshackle/core/utils/subprocess.py index ab5e57da..97598c68 100644 --- a/unshackle/core/utils/subprocess.py +++ b/unshackle/core/utils/subprocess.py @@ -145,7 +145,7 @@ def ffprobe(uri: Union[bytes, Path]) -> dict: if not binaries.FFProbe: raise EnvironmentError('FFProbe executable "ffprobe" not found but is required.') - args = [binaries.FFProbe, "-v", "quiet", "-of", "json", "-show_streams"] + args: list[Union[str, Path]] = [binaries.FFProbe, "-v", "quiet", "-of", "json", "-show_streams"] if isinstance(uri, Path): args.extend( ["-f", "lavfi", "-i", "movie={}[out+subcc]".format(str(uri).replace("\\", "/").replace(":", "\\\\:"))] diff --git a/unshackle/core/utils/template_formatter.py b/unshackle/core/utils/template_formatter.py index da63b6a2..63dc869c 100644 --- a/unshackle/core/utils/template_formatter.py +++ b/unshackle/core/utils/template_formatter.py @@ -27,7 +27,7 @@ class TemplateFormatter: Example: '{title}.{year}.{quality?}.{source}-{tag}' """ - def __init__(self, template: str, spacer: str | None = None): + def __init__(self, template: str, spacer: str | None = None) -> None: """Initialise the template formatter. Args: diff --git a/unshackle/core/vault.py b/unshackle/core/vault.py index 80ddc5fe..83b8b5bd 100644 --- a/unshackle/core/vault.py +++ b/unshackle/core/vault.py @@ -16,7 +16,7 @@ RATE_LIMIT_RETRIES = 3 class RateLimited(Exception): """The vault still answered 429 after the retries.""" - def __init__(self, message: str = f"still rate limited after {RATE_LIMIT_RETRIES} retries"): + def __init__(self, message: str = f"still rate limited after {RATE_LIMIT_RETRIES} retries") -> None: super().__init__(message) @@ -53,7 +53,7 @@ class Vault(metaclass=ABCMeta): local: bool = False report_now: bool = False - def __init__(self, name: str, no_push: bool = False): + def __init__(self, name: str, no_push: bool = False) -> None: self.name = name self.no_push = no_push diff --git a/unshackle/core/vaults.py b/unshackle/core/vaults.py index 92c4abb4..f60c432f 100644 --- a/unshackle/core/vaults.py +++ b/unshackle/core/vaults.py @@ -22,9 +22,9 @@ MODULES = {path.stem: getattr(import_module_by_path(path), path.stem) for path i class Vaults: """Keeps hold of Key Vaults with convenience functions, e.g. searching all vaults.""" - def __init__(self, service: Optional[str] = None): + def __init__(self, service: Optional[str] = None) -> None: self.service = service or "" - self.vaults = [] + self.vaults: list[Vault] = [] self.sources: dict[Union[UUID, str], tuple[str, Vault]] = {} self.flagged: set[tuple[Union[UUID, str], str]] = set() self.candidates: dict[Union[UUID, str], list[tuple[str, Vault]]] = {} diff --git a/unshackle/services/EXAMPLE/__init__.py b/unshackle/services/EXAMPLE/__init__.py index e1a6fd49..799bbe5c 100644 --- a/unshackle/services/EXAMPLE/__init__.py +++ b/unshackle/services/EXAMPLE/__init__.py @@ -110,7 +110,7 @@ class EXAMPLE(Service): def cli(ctx: click.Context, **kwargs: Any) -> EXAMPLE: return EXAMPLE(ctx, **kwargs) - def __init__(self, ctx: click.Context, title: str, movie: bool, device: str): + def __init__(self, ctx: click.Context, title: str, movie: bool, device: str) -> None: # Store CLI args BEFORE super().__init__ if the base needs them; here order # doesn't matter, but storing first is the common convention. self.title = title diff --git a/unshackle/services/MUSIC_EXAMPLE/__init__.py b/unshackle/services/MUSIC_EXAMPLE/__init__.py index 35186b54..e49b511f 100644 --- a/unshackle/services/MUSIC_EXAMPLE/__init__.py +++ b/unshackle/services/MUSIC_EXAMPLE/__init__.py @@ -102,7 +102,7 @@ class MUSIC_EXAMPLE(Service): def cli(ctx: click.Context, **kwargs: Any) -> MUSIC_EXAMPLE: return MUSIC_EXAMPLE(ctx, **kwargs) - def __init__(self, ctx: click.Context, title: str, device: str, lyrics: bool): + def __init__(self, ctx: click.Context, title: str, device: str, lyrics: bool) -> None: self.title = title self.device = device self.lyrics = lyrics diff --git a/unshackle/utils/base62.py b/unshackle/utils/base62.py index b4b4ad9e..9e1707de 100644 --- a/unshackle/utils/base62.py +++ b/unshackle/utils/base62.py @@ -15,10 +15,10 @@ CHARSET_DEFAULT = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxy CHARSET_INVERTED = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" -def encode(n, charset=CHARSET_DEFAULT): +def encode(n: int, charset: str = CHARSET_DEFAULT) -> str: """Encodes a given integer ``n``.""" - chs = [] + chs: list[str] = [] while n > 0: n, r = divmod(n, BASE) chs.insert(0, charset[r]) @@ -29,7 +29,7 @@ def encode(n, charset=CHARSET_DEFAULT): return "".join(chs) -def encodebytes(barray, charset=CHARSET_DEFAULT): +def encodebytes(barray: bytes, charset: str = CHARSET_DEFAULT) -> str: """Encodes a bytestring into a base62 string. :param barray: A byte array @@ -60,7 +60,7 @@ def encodebytes(barray, charset=CHARSET_DEFAULT): return zero_padding + value -def decode(encoded, charset=CHARSET_DEFAULT): +def decode(encoded: str, charset: str = CHARSET_DEFAULT) -> int: """Decodes a base62 encoded value ``encoded``. :type encoded: str @@ -76,7 +76,7 @@ def decode(encoded, charset=CHARSET_DEFAULT): return v -def decodebytes(encoded, charset=CHARSET_DEFAULT): +def decodebytes(encoded: str, charset: str = CHARSET_DEFAULT) -> bytes: """Decodes a string of base62 data into a bytes object. :param encoded: A string to be decoded in base62 @@ -99,7 +99,7 @@ def decodebytes(encoded, charset=CHARSET_DEFAULT): return leading_null_bytes + bytes(buf) -def value(ch, charset): +def value(ch: str, charset: str) -> int: """Decodes an individual digit of a base62 encoded string.""" try: @@ -108,7 +108,7 @@ def value(ch, charset): raise ValueError("base62: Invalid character (%s)" % ch) -def check_type(value, expected_type): +def check_type(value: object, expected_type: type) -> None: """Checks if the input is in an appropriate type.""" if not isinstance(value, expected_type): diff --git a/unshackle/vaults/API.py b/unshackle/vaults/API.py index e8e0765c..3c2b3f64 100644 --- a/unshackle/vaults/API.py +++ b/unshackle/vaults/API.py @@ -17,7 +17,7 @@ class API(Vault): no_push: bool = False, headers: Optional[dict[str, str]] = None, timeout: float = 10.0, - ): + ) -> None: super().__init__(name, no_push) self.uri = uri.rstrip("/") self.timeout = timeout diff --git a/unshackle/vaults/HTTP.py b/unshackle/vaults/HTTP.py index 820c89d0..b6c98a09 100644 --- a/unshackle/vaults/HTTP.py +++ b/unshackle/vaults/HTTP.py @@ -32,7 +32,7 @@ class HTTP(Vault): api_mode: str = "query", no_push: bool = False, timeout: float = 10.0, - ): + ) -> None: """ Initialise HTTP Vault. @@ -55,7 +55,7 @@ class HTTP(Vault): self.username = username self.api_mode = api_mode.lower() self.batch_insert = True - self.current_title = None + self.current_title: Optional[str] = None self.session = vault_session(timeout) self.api_session_id = None @@ -69,7 +69,7 @@ class HTTP(Vault): elif self.api_mode not in ["query", "json", "decrypt_labs"]: raise ValueError("api_mode must be either 'query', 'json', or 'decrypt_labs'") - def request(self, method: str, params: dict = None) -> dict: + def request(self, method: str, params: Optional[dict] = None) -> dict: """Make a request to the JSON API vault.""" if self.api_mode != "json": raise ValueError("request method is only available in json mode") @@ -165,10 +165,10 @@ class HTTP(Vault): "service": service.lower(), } - response = self.request("GetKey", params) - if response.get("status") == "not_found": + res = self.request("GetKey", params) + if res.get("status") == "not_found": return None - keys = response.get("keys", []) + keys = res.get("keys", []) for key_entry in keys: if isinstance(key_entry, str) and ":" in key_entry: entry_kid, entry_key = key_entry.split(":", 1) @@ -197,11 +197,11 @@ class HTTP(Vault): def get_keys(self, service: str) -> Iterator[tuple[str, str]]: if self.api_mode == "decrypt_labs": - return iter([]) + return elif self.api_mode == "json": # JSON API doesn't support getting all keys, so return empty iterator # This will cause the copy command to rely on the API's internal duplicate handling - return iter([]) + return else: response = self.session.get( self.url, @@ -231,7 +231,7 @@ class HTTP(Vault): if self.api_mode == "json": try: - response = self.request( + res = self.request( "InsertKey", { "kid": kid, @@ -240,9 +240,9 @@ class HTTP(Vault): "title": title, }, ) - if response.get("status") == "not_found": + if res.get("status") == "not_found": return False - return response.get("inserted", False) + return res.get("inserted", False) except Exception: return False else: @@ -282,18 +282,18 @@ class HTTP(Vault): if self.api_mode == "json": if self.batch_insert: try: - response = self.request( + res = self.request( "InsertKeys", {"keys": processed_kid_keys, "service": service.lower(), "title": title} ) - if response.get("status") != "not_found": - return int(response.get("inserted", 0)) + if res.get("status") != "not_found": + return int(res.get("inserted", 0)) self.batch_insert = False except ValueError as e: if "unknown" in str(e).lower() and "method" in str(e).lower(): self.batch_insert = False for kid, key in processed_kid_keys.items(): try: - response = self.request( + res = self.request( "InsertKey", { "kid": kid, @@ -302,9 +302,9 @@ class HTTP(Vault): "title": title, }, ) - if response.get("status") == "not_found": + if res.get("status") == "not_found": continue - if response.get("inserted", False): + if res.get("inserted", False): inserted_count += 1 except ValueError as e: failures.append(e) @@ -334,15 +334,15 @@ class HTTP(Vault): def get_services(self) -> Iterator[str]: if self.api_mode == "decrypt_labs": - return iter([]) + return elif self.api_mode == "json": try: - response = self.request("GetServices") - services = response.get("services", []) + res = self.request("GetServices") + services = res.get("services", []) for service in services: yield service except Exception: - return iter([]) + return else: response = self.session.get( self.url, @@ -359,7 +359,7 @@ class HTTP(Vault): for service in services: yield service - def set_title(self, title: str): + def set_title(self, title: str) -> None: """ Set a title for the next content key insertions. This is optional and will be sent with add_key requests if available. @@ -388,7 +388,7 @@ class HTTP(Vault): if self.api_mode == "json": try: - response = self.request( + res = self.request( "InsertKey", { "kid": kid, @@ -398,10 +398,10 @@ class HTTP(Vault): }, ) - if response.get("status") == "not_found": + if res.get("status") == "not_found": return InsertResult.FAILURE - if response.get("inserted", False): + if res.get("inserted", False): return InsertResult.SUCCESS else: return InsertResult.ALREADY_EXISTS diff --git a/unshackle/vaults/MySQL.py b/unshackle/vaults/MySQL.py index 8c8a0c52..2a2123c2 100644 --- a/unshackle/vaults/MySQL.py +++ b/unshackle/vaults/MySQL.py @@ -1,5 +1,7 @@ +from __future__ import annotations + import threading -from typing import Iterator, Optional, Union +from typing import Any, Iterator, Optional, Union from uuid import UUID import pymysql @@ -12,7 +14,9 @@ from unshackle.core.vault import Vault class MySQL(Vault): """Key Vault using a remotely-accessed mysql database connection.""" - def __init__(self, name: str, host: str, database: str, username: str, no_push: bool = False, **kwargs): + def __init__( + self, name: str, host: str, database: str, username: str, no_push: bool = False, **kwargs: Any + ) -> None: """ All extra arguments given through **kwargs will be sent to pymysql.connect. Use this to give more specific connection information. @@ -208,11 +212,13 @@ class MySQL(Vault): "SELECT count(TABLE_NAME) FROM information_schema.TABLES WHERE TABLE_SCHEMA=%s AND TABLE_NAME=%s", (conn.db, name), ) - return list(cursor.fetchone().values())[0] == 1 + row = cursor.fetchone() + assert row is not None # COUNT() always returns one row + return list(row.values())[0] == 1 finally: cursor.close() - def create_table(self, name: str): + def create_table(self, name: str) -> None: """Make a Table with the specified name if it does not exist yet.""" if self.has_table(name): return @@ -246,8 +252,7 @@ class MySQL(Vault): try: cursor.execute("SHOW GRANTS") - grants = cursor.fetchall() - grants = [next(iter(x.values())) for x in grants] + grants = [next(iter(x.values())) for x in cursor.fetchall()] grants = [tuple(x[6:].split(" TO ")[0].split(" ON ")) for x in list(grants)] grants = [ ( @@ -272,14 +277,14 @@ class MySQL(Vault): class ConnectionFactory: - def __init__(self, con: dict): + def __init__(self, con: dict) -> None: self._con = con self._store = threading.local() - def create_connection(self) -> pymysql.Connection: + def create_connection(self) -> pymysql.Connection[DictCursor]: return pymysql.connect(**self._con) - def get(self) -> pymysql.Connection: + def get(self) -> pymysql.Connection[DictCursor]: if not hasattr(self._store, "conn"): self._store.conn = self.create_connection() return self._store.conn diff --git a/unshackle/vaults/SQLite.py b/unshackle/vaults/SQLite.py index b2219c66..b8bdac54 100644 --- a/unshackle/vaults/SQLite.py +++ b/unshackle/vaults/SQLite.py @@ -13,7 +13,7 @@ class SQLite(Vault): local = True - def __init__(self, name: str, path: Union[str, Path], no_push: bool = False): + def __init__(self, name: str, path: Union[str, Path], no_push: bool = False) -> None: super().__init__(name, no_push) self.path = Path(path).expanduser() # TODO: Use a DictCursor or such to get fetches as dict? @@ -113,7 +113,7 @@ class SQLite(Vault): self.create_bad_keys_table() bad = set(self.conn_factory.get().execute("SELECT lower(`kid`), lower(`key_`) FROM `bad_keys`").fetchall()) kid_keys = { - kid: key_ for kid, key_ in kid_keys.items() if (kid.replace("-", "").lower(), key_.lower()) not in bad + kid: key_ for kid, key_ in kid_keys.items() if (str(kid).replace("-", "").lower(), key_.lower()) not in bad } if not kid_keys: @@ -258,7 +258,7 @@ class SQLite(Vault): finally: cursor.close() - def create_table(self, name: str): + def create_table(self, name: str) -> None: """Make a Table with the specified name if it does not exist yet.""" if self.has_table(name): return @@ -285,7 +285,7 @@ class SQLite(Vault): class ConnectionFactory: - def __init__(self, path: Union[str, Path]): + def __init__(self, path: Union[str, Path]) -> None: self._path = path self._store = threading.local() diff --git a/uv.lock b/uv.lock index f112c7a9..7e3a9b88 100644 --- a/uv.lock +++ b/uv.lock @@ -2129,6 +2129,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8b/fa/4f4d3bfca9ef6dd17d69ed18b96564c53b32d3ce774132308d0bee849f10/types_pymysql-1.1.0.20251220-py3-none-any.whl", hash = "sha256:fa1082af7dea6c53b6caa5784241924b1296ea3a8d3bd060417352c5e10c0618", size = 23067, upload-time = "2025-12-20T03:07:37.766Z" }, ] +[[package]] +name = "types-pyyaml" +version = "6.0.12.20261030" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c7/0c/72b85bc2e3a3cf23419a7f8861c1855893f02080e202b631e33cc0d1a94a/types_pyyaml-6.0.12.20261030.tar.gz", hash = "sha256:9123d707aa31a3883eeb2cb290537c75246dfcb1c00abed7a70f5419d3bd31ae", size = 18878, upload-time = "2026-10-30T07:50:19.43Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/34/99/7e2c7c7202108220fb664ce4c5ac313a564ab5babb153516a3ac0ad8a681/types_pyyaml-6.0.12.20261030-py3-none-any.whl", hash = "sha256:508cc80546918201026a955356222111224e1f4f7deab0b91f2a95933f764e67", size = 21283, upload-time = "2026-10-30T07:50:18.389Z" }, +] + [[package]] name = "types-requests" version = "2.32.4.20260107" @@ -2220,6 +2229,7 @@ dev = [ { name = "ruff" }, { name = "types-protobuf" }, { name = "types-pymysql" }, + { name = "types-pyyaml" }, { name = "types-requests" }, { name = "virtualenv" }, ] @@ -2290,6 +2300,7 @@ dev = [ { name = "ruff", specifier = ">=0.15.12,<0.16" }, { name = "types-protobuf", specifier = ">=4.24.0.20240408,<7" }, { name = "types-pymysql", specifier = ">=1.1.0.1,<2" }, + { name = "types-pyyaml", specifier = ">=6.0,<7" }, { name = "types-requests", specifier = ">=2.31.0.20240406,<3" }, { name = "virtualenv", specifier = ">=20.36.1,<22" }, ]