From fa3cee11b7e6f701d5fcfb900986531f1df13996 Mon Sep 17 00:00:00 2001 From: rlaphoenix Date: Tue, 9 Jan 2024 11:41:57 +0000 Subject: [PATCH] Move Download Cancel/Skip Events to constants --- devine/commands/dl.py | 32 ++++++++++++------------------ devine/core/constants.py | 4 ++++ devine/core/manifests/dash.py | 29 +++++++++++---------------- devine/core/manifests/hls.py | 37 ++++++++++++----------------------- 4 files changed, 41 insertions(+), 61 deletions(-) diff --git a/devine/commands/dl.py b/devine/commands/dl.py index 32cea3f..30eca40 100644 --- a/devine/commands/dl.py +++ b/devine/commands/dl.py @@ -16,7 +16,7 @@ from functools import partial from http.cookiejar import MozillaCookieJar from itertools import zip_longest from pathlib import Path -from threading import Event, Lock +from threading import Lock from typing import Any, Callable, Optional from uuid import UUID @@ -41,7 +41,7 @@ from rich.tree import Tree from devine.core.config import config from devine.core.console import console -from devine.core.constants import AnyTrack, context_settings +from devine.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack, context_settings from devine.core.credential import Credential from devine.core.downloaders import downloader from devine.core.drm import DRM_T, Widevine @@ -138,8 +138,6 @@ class dl: def cli(ctx: click.Context, **kwargs: Any) -> dl: return dl(ctx, **kwargs) - DL_POOL_STOP = Event() - DL_POOL_SKIP = Event() DRM_TABLE_LOCK = Lock() def __init__( @@ -466,7 +464,7 @@ class dl: dl_start_time = time.time() if skip_dl: - self.DL_POOL_SKIP.set() + DOWNLOAD_LICENCE_ONLY.set() try: with Live( @@ -827,11 +825,11 @@ class dl: prepare_drm: Callable, progress: partial ): - if self.DL_POOL_SKIP.is_set(): + if DOWNLOAD_LICENCE_ONLY.is_set(): progress(downloaded="[yellow]SKIPPING") - if self.DL_POOL_STOP.is_set(): - progress(downloaded="[yellow]SKIPPED") + if DOWNLOAD_CANCELLED.is_set(): + progress(downloaded="[yellow]CANCELLED") return proxy = next(iter(service.session.proxies.values()), None) @@ -856,7 +854,7 @@ class dl: if save_dir.exists() and save_dir.name.endswith("_segments"): shutil.rmtree(save_dir) - if not self.DL_POOL_SKIP.is_set(): + if not DOWNLOAD_LICENCE_ONLY.is_set(): if config.directories.temp.is_file(): self.log.error(f"Temp Directory '{config.directories.temp}' must be a Directory, not a file") sys.exit(1) @@ -875,8 +873,6 @@ class dl: track=track, save_path=save_path, save_dir=save_dir, - stop_event=self.DL_POOL_STOP, - skip_event=self.DL_POOL_SKIP, progress=progress, session=service.session, proxy=proxy, @@ -887,8 +883,6 @@ class dl: track=track, save_path=save_path, save_dir=save_dir, - stop_event=self.DL_POOL_STOP, - skip_event=self.DL_POOL_SKIP, progress=progress, session=service.session, proxy=proxy, @@ -919,7 +913,7 @@ class dl: else: drm = None - if self.DL_POOL_SKIP.is_set(): + if DOWNLOAD_LICENCE_ONLY.is_set(): progress(downloaded="[yellow]SKIPPED") else: downloader( @@ -948,23 +942,23 @@ class dl: progress(downloaded="Downloaded") except KeyboardInterrupt: - self.DL_POOL_STOP.set() + DOWNLOAD_CANCELLED.set() progress(downloaded="[yellow]CANCELLED") raise except Exception: - self.DL_POOL_STOP.set() + DOWNLOAD_CANCELLED.set() progress(downloaded="[red]FAILED") raise except (Exception, KeyboardInterrupt): - if not self.DL_POOL_SKIP.is_set(): + if not DOWNLOAD_LICENCE_ONLY.is_set(): cleanup() raise - if self.DL_POOL_STOP.is_set(): + if DOWNLOAD_CANCELLED.is_set(): # we stopped during the download, let's exit return - if not self.DL_POOL_SKIP.is_set(): + if not DOWNLOAD_LICENCE_ONLY.is_set(): if track.path.stat().st_size <= 3: # Empty UTF-8 BOM == 3 bytes raise IOError("Download failed, the downloaded file is empty.") diff --git a/devine/core/constants.py b/devine/core/constants.py index 4e6efa4..447351c 100644 --- a/devine/core/constants.py +++ b/devine/core/constants.py @@ -1,5 +1,9 @@ +from threading import Event from typing import TypeVar, Union +DOWNLOAD_CANCELLED = Event() +DOWNLOAD_LICENCE_ONLY = Event() + DRM_SORT_MAP = ["ClearKey", "Widevine"] LANGUAGE_MUX_MAP = { # List of language tags that cannot be used by mkvmerge and need replacements. diff --git a/devine/core/manifests/dash.py b/devine/core/manifests/dash.py index b0e1f08..4fc812c 100644 --- a/devine/core/manifests/dash.py +++ b/devine/core/manifests/dash.py @@ -12,7 +12,6 @@ from copy import copy from functools import partial from hashlib import md5 from pathlib import Path -from threading import Event from typing import Any, Callable, MutableMapping, Optional, Union from urllib.parse import urljoin, urlparse from uuid import UUID @@ -26,7 +25,7 @@ from requests import Session from requests.cookies import RequestsCookieJar from rich import filesize -from devine.core.constants import AnyTrack +from devine.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack from devine.core.downloaders import downloader from devine.core.downloaders import requests as requests_downloader from devine.core.drm import Widevine @@ -225,8 +224,6 @@ class DASH: track: AnyTrack, save_path: Path, save_dir: Path, - stop_event: Event, - skip_event: Event, progress: partial, session: Optional[Session] = None, proxy: Optional[str] = None, @@ -401,13 +398,13 @@ class DASH: license_widevine(drm, track_kid=track_kid) progress(downloaded="[yellow]LICENSED") except Exception: # noqa - stop_event.set() # skip pending track downloads + DOWNLOAD_CANCELLED.set() # skip pending track downloads progress(downloaded="[red]FAILED") raise else: drm = None - if skip_event.is_set(): + if DOWNLOAD_LICENCE_ONLY.is_set(): progress(downloaded="[yellow]SKIPPED") return @@ -427,15 +424,14 @@ class DASH: proxy=proxy, headers=session.headers, cookies=session.cookies, - bytes_range=bytes_range, - stop_event=stop_event + bytes_range=bytes_range ) for n, (url, bytes_range) in enumerate(segments) ))): try: download_size = download.result() except KeyboardInterrupt: - stop_event.set() # skip pending track downloads + DOWNLOAD_CANCELLED.set() # skip pending track downloads progress(downloaded="[yellow]CANCELLING") pool.shutdown(wait=True, cancel_futures=True) progress(downloaded="[yellow]CANCELLED") @@ -443,7 +439,7 @@ class DASH: # the pool is already shut down, so exiting loop is fine raise except Exception: - stop_event.set() # skip pending track downloads + DOWNLOAD_CANCELLED.set() # skip pending track downloads progress(downloaded="[red]FAILING") pool.shutdown(wait=True, cancel_futures=True) progress(downloaded="[red]FAILED") @@ -501,8 +497,7 @@ class DASH: proxy: Optional[str] = None, headers: Optional[MutableMapping[str, str | bytes]] = None, cookies: Optional[Union[MutableMapping[str, str], RequestsCookieJar]] = None, - bytes_range: Optional[str] = None, - stop_event: Optional[Event] = None + bytes_range: Optional[str] = None ) -> int: """ Download a DASH Media Segment. @@ -518,12 +513,10 @@ class DASH: will be resolved based on the URI among other parameters. Multiple cookies with the same name but a different domain/path are resolved. bytes_range: Download only specific bytes of the Segment file using the Range header. - stop_event: Prematurely stop the Download from beginning. Useful if ran from - a Thread Pool. It will raise a KeyboardInterrupt if set. Returns the file size of the downloaded Segment in bytes. """ - if stop_event and stop_event.is_set(): + if DOWNLOAD_CANCELLED.is_set(): raise KeyboardInterrupt() attempts = 1 @@ -546,9 +539,9 @@ class DASH: segmented=True ) break - except Exception as ee: - if (stop_event and stop_event.is_set()) or attempts == 5: - raise ee + except Exception as e: + if DOWNLOAD_CANCELLED.is_set() or attempts == 5: + raise e time.sleep(2) attempts += 1 diff --git a/devine/core/manifests/hls.py b/devine/core/manifests/hls.py index 93060f5..1e3a775 100644 --- a/devine/core/manifests/hls.py +++ b/devine/core/manifests/hls.py @@ -10,7 +10,7 @@ from functools import partial from hashlib import md5 from pathlib import Path from queue import Queue -from threading import Event, Lock +from threading import Lock from typing import Any, Callable, Optional, Union from urllib.parse import urljoin @@ -23,7 +23,7 @@ from pywidevine.pssh import PSSH from requests import Session from rich import filesize -from devine.core.constants import AnyTrack +from devine.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack from devine.core.downloaders import downloader from devine.core.downloaders import requests as requests_downloader from devine.core.drm import DRM_T, ClearKey, Widevine @@ -190,8 +190,6 @@ class HLS: track: AnyTrack, save_path: Path, save_dir: Path, - stop_event: Event, - skip_event: Event, progress: partial, session: Optional[Session] = None, proxy: Optional[str] = None, @@ -231,7 +229,7 @@ class HLS: license_widevine(session_drm) progress(downloaded="[yellow]LICENSED") except Exception: # noqa - stop_event.set() # skip pending track downloads + DOWNLOAD_CANCELLED.set() # skip pending track downloads progress(downloaded="[red]FAILED") raise else: @@ -265,16 +263,14 @@ class HLS: progress=progress, license_widevine=license_widevine, session=session, - proxy=proxy, - stop_event=stop_event, - skip_event=skip_event + proxy=proxy ) for n, segment in enumerate(master.segments) ))): try: download_size = download.result() except KeyboardInterrupt: - stop_event.set() # skip pending track downloads + DOWNLOAD_CANCELLED.set() # skip pending track downloads progress(downloaded="[yellow]CANCELLING") pool.shutdown(wait=True, cancel_futures=True) progress(downloaded="[yellow]CANCELLED") @@ -282,7 +278,7 @@ class HLS: # the pool is already shut down, so exiting loop is fine raise except Exception as e: - stop_event.set() # skip pending track downloads + DOWNLOAD_CANCELLED.set() # skip pending track downloads progress(downloaded="[red]FAILING") pool.shutdown(wait=True, cancel_futures=True) progress(downloaded="[red]FAILED") @@ -310,7 +306,7 @@ class HLS: last_speed_refresh = now download_sizes.clear() - if skip_event.is_set(): + if DOWNLOAD_LICENCE_ONLY.is_set(): return with open(save_path, "wb") as f: @@ -338,9 +334,7 @@ class HLS: progress: partial, license_widevine: Optional[Callable] = None, session: Optional[Session] = None, - proxy: Optional[str] = None, - stop_event: Optional[Event] = None, - skip_event: Optional[Event] = None + proxy: Optional[str] = None ) -> int: """ Download (and Decrypt) an HLS Media Segment. @@ -365,15 +359,10 @@ class HLS: if the Segment's DRM uses Widevine. proxy: Proxy URI to use when downloading the Segment file. session: Python-Requests Session used when requesting init data. - stop_event: Prematurely stop the Download from beginning. Useful if ran from - a Thread Pool. It will raise a KeyboardInterrupt if set. - skip_event: Prematurely stop the Download from beginning. It returns with a - file size of -1 directly after DRM licensing occurs, even if it's DRM-free. - This is mainly for `--skip-dl` to allow licensing without downloading. Returns the file size of the downloaded Segment in bytes. """ - if stop_event.is_set(): + if DOWNLOAD_CANCELLED.is_set(): raise KeyboardInterrupt() if callable(track.OnSegmentFilter) and track.OnSegmentFilter(segment): @@ -430,7 +419,7 @@ class HLS: finally: segment_key.put(newest_segment_key) - if skip_event.is_set(): + if DOWNLOAD_LICENCE_ONLY.is_set(): return -1 attempts = 1 @@ -456,9 +445,9 @@ class HLS: segmented=True ) break - except Exception as ee: - if stop_event.is_set() or attempts == 5: - raise ee + except Exception as e: + if DOWNLOAD_CANCELLED.is_set() or attempts == 5: + raise e time.sleep(2) attempts += 1