refactor: Move dl command's download_track() to Track.download()

devine/commands/dl.py

@@ -40,18 +40,16 @@ from rich.tree import Tree
 
 from devine.core.config import config
 from devine.core.console import console
-from devine.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack, context_settings
+from devine.core.constants import 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
-from devine.core.manifests import DASH, HLS
 from devine.core.proxies import Basic, Hola, NordVPN
 from devine.core.service import Service
 from devine.core.services import Services
 from devine.core.titles import Movie, Song, Title_T
 from devine.core.titles.episode import Episode
 from devine.core.tracks import Audio, Subtitle, Tracks, Video
-from devine.core.utilities import get_binary_path, is_close_match, time_elapsed_since, try_ensure_utf8
+from devine.core.utilities import get_binary_path, is_close_match, time_elapsed_since
 from devine.core.utils.click_types import LANGUAGE_RANGE, QUALITY_LIST, SEASON_RANGE, ContextData
 from devine.core.utils.collections import merge_dict
 from devine.core.utils.subprocess import ffprobe
@@ -376,13 +374,7 @@ class dl:
                             sys.exit(1)
 
                     video_languages = v_lang or lang
-                    if (
+                    if video_languages and "all" not in video_languages:
-                        (v_lang and "all" not in v_lang) or
-                        (lang and "all" not in lang and any(
-                            x.language != title.tracks.videos[0].language
-                            for x in title.tracks.videos
-                        ))
-                    ):
                         title.tracks.videos = title.tracks.by_language(title.tracks.videos, video_languages)
                         if not title.tracks.videos:
                             self.log.error(f"There's no {video_languages} Video Track...")
@@ -476,9 +468,8 @@ class dl:
                     with ThreadPoolExecutor(workers) as pool:
                         for download in futures.as_completed((
                             pool.submit(
-                                self.download_track,
+                                track.download,
-                                service=service,
+                                session=service.session,
-                                track=track,
                                 prepare_drm=partial(
                                     partial(
                                         self.prepare_drm,
@@ -795,156 +786,6 @@ class dl:
                     keys[str(title)][str(track)].update(drm.content_keys)
                     export.write_text(jsonpickle.dumps(keys, indent=4), encoding="utf8")
 
-    def download_track(
-        self,
-        service: Service,
-        track: AnyTrack,
-        prepare_drm: Callable,
-        progress: partial
-    ):
-        if DOWNLOAD_LICENCE_ONLY.is_set():
-            progress(downloaded="[yellow]SKIPPING")
-
-        if DOWNLOAD_CANCELLED.is_set():
-            progress(downloaded="[yellow]CANCELLED")
-            return
-
-        proxy = next(iter(service.session.proxies.values()), None)
-
-        save_path = config.directories.temp / f"{track.__class__.__name__}_{track.id}.mp4"
-        if isinstance(track, Subtitle):
-            save_path = save_path.with_suffix(f".{track.codec.extension}")
-
-        if track.descriptor != track.Descriptor.URL:
-            save_dir = save_path.with_name(save_path.name + "_segments")
-        else:
-            save_dir = save_path.parent
-
-        def cleanup():
-            # track file (e.g., "foo.mp4")
-            save_path.unlink(missing_ok=True)
-            # aria2c control file (e.g., "foo.mp4.aria2")
-            save_path.with_suffix(f"{save_path.suffix}.aria2").unlink(missing_ok=True)
-            if save_dir.exists() and save_dir.name.endswith("_segments"):
-                shutil.rmtree(save_dir)
-
-        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)
-
-            config.directories.temp.mkdir(parents=True, exist_ok=True)
-
-            # Delete any pre-existing temp files matching this track.
-            # We can't re-use or continue downloading these tracks as they do not use a
-            # lock file. Or at least the majority don't. Even if they did I've encountered
-            # corruptions caused by sudden interruptions to the lock file.
-            cleanup()
-
-        try:
-            if track.descriptor == track.Descriptor.HLS:
-                HLS.download_track(
-                    track=track,
-                    save_path=save_path,
-                    save_dir=save_dir,
-                    progress=progress,
-                    session=service.session,
-                    proxy=proxy,
-                    license_widevine=prepare_drm
-                )
-            elif track.descriptor == track.Descriptor.DASH:
-                DASH.download_track(
-                    track=track,
-                    save_path=save_path,
-                    save_dir=save_dir,
-                    progress=progress,
-                    session=service.session,
-                    proxy=proxy,
-                    license_widevine=prepare_drm
-                )
-            elif track.descriptor == track.Descriptor.URL:
-                try:
-                    if not track.drm and isinstance(track, (Video, Audio)):
-                        # the service might not have explicitly defined the `drm` property
-                        # try find widevine DRM information from the init data of URL
-                        try:
-                            track.drm = [Widevine.from_track(track, service.session)]
-                        except Widevine.Exceptions.PSSHNotFound:
-                            # it might not have Widevine DRM, or might not have found the PSSH
-                            self.log.warning("No Widevine PSSH was found for this track, is it DRM free?")
-
-                    if track.drm:
-                        track_kid = track.get_key_id(session=service.session)
-                        drm = track.drm[0]  # just use the first supported DRM system for now
-                        if isinstance(drm, Widevine):
-                            # license and grab content keys
-                            if not prepare_drm:
-                                raise ValueError("prepare_drm func must be supplied to use Widevine DRM")
-                            progress(downloaded="LICENSING")
-                            prepare_drm(drm, track_kid=track_kid)
-                            progress(downloaded="[yellow]LICENSED")
-                    else:
-                        drm = None
-
-                    if DOWNLOAD_LICENCE_ONLY.is_set():
-                        progress(downloaded="[yellow]SKIPPED")
-                    else:
-                        for status_update in downloader(
-                            urls=track.url,
-                            output_dir=save_path.parent,
-                            filename=save_path.name,
-                            headers=service.session.headers,
-                            cookies=service.session.cookies,
-                            proxy=proxy
-                        ):
-                            file_downloaded = status_update.get("file_downloaded")
-                            if not file_downloaded:
-                                progress(**status_update)
-
-                        track.path = save_path
-                        if callable(track.OnDownloaded):
-                            track.OnDownloaded()
-
-                        if drm:
-                            progress(downloaded="Decrypting", completed=0, total=100)
-                            drm.decrypt(save_path)
-                            track.drm = None
-                            if callable(track.OnDecrypted):
-                                track.OnDecrypted(drm)
-                            progress(downloaded="Decrypted", completed=100)
-
-                        if isinstance(track, Subtitle) and \
-                           track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML):
-                            track_data = track.path.read_bytes()
-                            track_data = try_ensure_utf8(track_data)
-                            track_data = track_data.decode("utf8"). \
-                                replace("‎", html.unescape("‎")). \
-                                replace("‏", html.unescape("‏")). \
-                                encode("utf8")
-                            track.path.write_bytes(track_data)
-
-                        progress(downloaded="Downloaded")
-                except KeyboardInterrupt:
-                    DOWNLOAD_CANCELLED.set()
-                    progress(downloaded="[yellow]CANCELLED")
-                    raise
-                except Exception:
-                    DOWNLOAD_CANCELLED.set()
-                    progress(downloaded="[red]FAILED")
-                    raise
-        except (Exception, KeyboardInterrupt):
-            if not DOWNLOAD_LICENCE_ONLY.is_set():
-                cleanup()
-            raise
-
-        if DOWNLOAD_CANCELLED.is_set():
-            # we stopped during the download, let's exit
-            return
-
-        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.")
-
     @staticmethod
     def get_profile(service: str) -> Optional[str]:
         """Get profile for Service from config."""

devine/core/tracks/track.py

@@ -1,9 +1,12 @@
 import base64
+import html
+import logging
 import re
 import shutil
 import subprocess
 from copy import copy
 from enum import Enum
+from functools import partial
 from pathlib import Path
 from typing import Any, Callable, Iterable, Optional, Union
 from uuid import UUID
@@ -13,9 +16,11 @@ import m3u8
 import requests
 from langcodes import Language
 
-from devine.core.constants import TERRITORY_MAP
+from devine.core.config import config
-from devine.core.drm import DRM_T
+from devine.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, TERRITORY_MAP
-from devine.core.utilities import get_binary_path, get_boxes
+from devine.core.downloaders import downloader
+from devine.core.drm import DRM_T, Widevine
+from devine.core.utilities import get_binary_path, get_boxes, try_ensure_utf8
 from devine.core.utils.subprocess import ffprobe
 
 
@@ -109,6 +114,163 @@ class Track:
     def __eq__(self, other: Any) -> bool:
         return isinstance(other, Track) and self.id == other.id
 
+    def download(
+        self,
+        session: requests.Session,
+        prepare_drm: partial,
+        progress: Optional[partial] = None
+    ):
+        """Download and optionally Decrypt this Track."""
+        from devine.core.manifests import DASH, HLS
+
+        if DOWNLOAD_LICENCE_ONLY.is_set():
+            progress(downloaded="[yellow]SKIPPING")
+
+        if DOWNLOAD_CANCELLED.is_set():
+            progress(downloaded="[yellow]SKIPPED")
+            return
+
+        log = logging.getLogger("track")
+
+        proxy = next(iter(session.proxies.values()), None)
+
+        track_type = self.__class__.__name__
+        save_path = config.directories.temp / f"{track_type}_{self.id}.mp4"
+        if track_type == "Subtitle":
+            save_path = save_path.with_suffix(f".{self.codec.extension}")
+
+        if self.descriptor != self.Descriptor.URL:
+            save_dir = save_path.with_name(save_path.name + "_segments")
+        else:
+            save_dir = save_path.parent
+
+        def cleanup():
+            # track file (e.g., "foo.mp4")
+            save_path.unlink(missing_ok=True)
+            # aria2c control file (e.g., "foo.mp4.aria2" or "foo.mp4.aria2__temp")
+            save_path.with_suffix(f"{save_path.suffix}.aria2").unlink(missing_ok=True)
+            save_path.with_suffix(f"{save_path.suffix}.aria2__temp").unlink(missing_ok=True)
+            if save_dir.exists() and save_dir.name.endswith("_segments"):
+                shutil.rmtree(save_dir)
+
+        if not DOWNLOAD_LICENCE_ONLY.is_set():
+            if config.directories.temp.is_file():
+                raise ValueError(f"Temp Directory '{config.directories.temp}' must be a Directory, not a file")
+
+            config.directories.temp.mkdir(parents=True, exist_ok=True)
+
+            # Delete any pre-existing temp files matching this track.
+            # We can't re-use or continue downloading these tracks as they do not use a
+            # lock file. Or at least the majority don't. Even if they did I've encountered
+            # corruptions caused by sudden interruptions to the lock file.
+            cleanup()
+
+        try:
+            if self.descriptor == self.Descriptor.HLS:
+                HLS.download_track(
+                    track=self,
+                    save_path=save_path,
+                    save_dir=save_dir,
+                    progress=progress,
+                    session=session,
+                    proxy=proxy,
+                    license_widevine=prepare_drm
+                )
+            elif self.descriptor == self.Descriptor.DASH:
+                DASH.download_track(
+                    track=self,
+                    save_path=save_path,
+                    save_dir=save_dir,
+                    progress=progress,
+                    session=session,
+                    proxy=proxy,
+                    license_widevine=prepare_drm
+                )
+            elif self.descriptor == self.Descriptor.URL:
+                try:
+                    if not self.drm and track_type in ("Video", "Audio"):
+                        # the service might not have explicitly defined the `drm` property
+                        # try find widevine DRM information from the init data of URL
+                        try:
+                            self.drm = [Widevine.from_track(self, session)]
+                        except Widevine.Exceptions.PSSHNotFound:
+                            # it might not have Widevine DRM, or might not have found the PSSH
+                            log.warning("No Widevine PSSH was found for this track, is it DRM free?")
+
+                    if self.drm:
+                        track_kid = self.get_key_id(session=session)
+                        drm = self.drm[0]  # just use the first supported DRM system for now
+                        if isinstance(drm, Widevine):
+                            # license and grab content keys
+                            if not prepare_drm:
+                                raise ValueError("prepare_drm func must be supplied to use Widevine DRM")
+                            progress(downloaded="LICENSING")
+                            prepare_drm(drm, track_kid=track_kid)
+                            progress(downloaded="[yellow]LICENSED")
+                    else:
+                        drm = None
+
+                    if DOWNLOAD_LICENCE_ONLY.is_set():
+                        progress(downloaded="[yellow]SKIPPED")
+                    else:
+                        for status_update in downloader(
+                            urls=self.url,
+                            output_dir=save_path.parent,
+                            filename=save_path.name,
+                            headers=session.headers,
+                            cookies=session.cookies,
+                            proxy=proxy
+                        ):
+                            file_downloaded = status_update.get("file_downloaded")
+                            if not file_downloaded:
+                                progress(**status_update)
+
+                        self.path = save_path
+                        if callable(self.OnDownloaded):
+                            self.OnDownloaded()
+
+                        if drm:
+                            progress(downloaded="Decrypting", completed=0, total=100)
+                            drm.decrypt(save_path)
+                            self.drm = None
+                            if callable(self.OnDecrypted):
+                                self.OnDecrypted(drm)
+                            progress(downloaded="Decrypted", completed=100)
+
+                        if track_type == "Subtitle" and self.codec.name not in ("fVTT", "fTTML"):
+                            track_data = self.path.read_bytes()
+                            track_data = try_ensure_utf8(track_data)
+                            track_data = track_data.decode("utf8"). \
+                                replace("‎", html.unescape("‎")). \
+                                replace("‏", html.unescape("‏")). \
+                                encode("utf8")
+                            self.path.write_bytes(track_data)
+
+                        progress(downloaded="Downloaded")
+                except KeyboardInterrupt:
+                    DOWNLOAD_CANCELLED.set()
+                    progress(downloaded="[yellow]CANCELLED")
+                    raise
+                except Exception:
+                    DOWNLOAD_CANCELLED.set()
+                    progress(downloaded="[red]FAILED")
+                    raise
+        except (Exception, KeyboardInterrupt):
+            if not DOWNLOAD_LICENCE_ONLY.is_set():
+                cleanup()
+            raise
+
+        if DOWNLOAD_CANCELLED.is_set():
+            # we stopped during the download, let's exit
+            return
+
+        if not DOWNLOAD_LICENCE_ONLY.is_set():
+            if self.path.stat().st_size <= 3:  # Empty UTF-8 BOM == 3 bytes
+                raise IOError("Download failed, the downloaded file is empty.")
+
+        if callable(self.OnDownloaded):
+            self.OnDownloaded(self)
+
     def delete(self) -> None:
         if self.path:
             self.path.unlink()