Move dl command's download_track code 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
@@ -477,9 +475,8 @@ class dl:
                     with ThreadPoolExecutor(workers) as pool:
                         for download in futures.as_completed((
                             pool.submit(
-                                self.download_track,
+                                track.download,
                                 service=service,
-                                track=track,
                                 prepare_drm=partial(
                                     partial(
                                         self.prepare_drm,
@@ -796,151 +793,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.M3U:
-                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.MPD:
-                DASH.download_track(
-                    track=track,
-                    save_path=save_path,
-                    save_dir=save_dir,
-                    progress=progress,
-                    session=service.session,
-                    proxy=proxy,
-                    license_widevine=prepare_drm
-                )
-            # no else-if as DASH may convert the track to URL descriptor
-            if 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:
-                        downloader(
-                            uri=track.url,
-                            out=save_path,
-                            headers=service.session.headers,
-                            cookies=service.session.cookies,
-                            proxy=proxy,
-                            progress=progress
-                        )
-
-                        track.path = save_path
-
-                        if drm:
-                            progress(downloaded="Decrypting", completed=0, total=100)
-                            drm.decrypt(save_path)
-                            track.drm = None
-                            if callable(track.OnDecrypted):
-                                track.OnDecrypted(track)
-                            progress(downloaded="Decrypted", completed=100)
-
-                        if isinstance(track, Subtitle):
-                            track_data = track.path.read_bytes()
-                            track_data = try_ensure_utf8(track_data)
-                            track_data = html.unescape(track_data.decode("utf8")).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.")
-
-            if callable(track.OnDownloaded):
-                track.OnDownloaded(track)
-
     @staticmethod
     def get_profile(service: str) -> Optional[str]:
         """Get profile for Service from config."""

devine/core/tracks/track.py

@@ -1,8 +1,11 @@
 import base64
+import html
+import logging
 import re
 import shutil
 import subprocess
 from enum import Enum
+from functools import partial
 from pathlib import Path
 from typing import Any, Callable, Iterable, Optional, Union
 from uuid import UUID
@@ -10,9 +13,12 @@ from uuid import UUID
 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.service import Service
+from devine.core.utilities import get_binary_path, get_boxes, try_ensure_utf8
 from devine.core.utils.subprocess import ffprobe
 
 
@@ -216,6 +222,148 @@ class Track:
 
         return init_data
 
+    def download(self, service: Service, prepare_drm: Callable, progress: partial) -> None:
+        if DOWNLOAD_LICENCE_ONLY.is_set():
+            progress(downloaded="[yellow]SKIPPING")
+
+        if DOWNLOAD_CANCELLED.is_set():
+            progress(downloaded="[yellow]CANCELLED")
+            return
+
+        log = logging.getLogger("Track")
+
+        proxy = next(iter(service.session.proxies.values()), None)
+
+        save_path = config.directories.temp / f"{self.__class__.__name__}_{self.id}.mp4"
+        if self.__class__.__name__ == "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")
+            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():
+                raise EnvironmentError(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.M3U:
+                from devine.core.manifests import HLS
+                HLS.download_track(
+                    track=self,
+                    save_path=save_path,
+                    save_dir=save_dir,
+                    progress=progress,
+                    session=service.session,
+                    proxy=proxy,
+                    license_widevine=prepare_drm
+                )
+            elif self.descriptor == self.Descriptor.MPD:
+                from devine.core.manifests import DASH
+                DASH.download_track(
+                    track=self,
+                    save_path=save_path,
+                    save_dir=save_dir,
+                    progress=progress,
+                    session=service.session,
+                    proxy=proxy,
+                    license_widevine=prepare_drm
+                )
+            # no else-if as DASH may convert the track to URL descriptor
+            if self.descriptor == self.Descriptor.URL:
+                try:
+                    if not self.drm and self.__class__.__name__ 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, service.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=service.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:
+                        downloader(
+                            uri=self.url,
+                            out=save_path,
+                            headers=service.session.headers,
+                            cookies=service.session.cookies,
+                            proxy=proxy,
+                            progress=progress
+                        )
+
+                        self.path = save_path
+
+                        if drm:
+                            progress(downloaded="Decrypting", completed=0, total=100)
+                            drm.decrypt(save_path)
+                            self.drm = None
+                            if callable(self.OnDecrypted):
+                                self.OnDecrypted(self)
+                            progress(downloaded="Decrypted", completed=100)
+
+                        if self.__class__.__name__ == "Subtitle":
+                            track_data = self.path.read_bytes()
+                            track_data = try_ensure_utf8(track_data)
+                            track_data = html.unescape(track_data.decode("utf8")).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()