diff --git a/devine/commands/dl.py b/devine/commands/dl.py index a760ad9..5af9eea 100644 --- a/devine/commands/dl.py +++ b/devine/commands/dl.py @@ -735,6 +735,7 @@ class dl: if track.descriptor == track.Descriptor.M3U: HLS.download_track( track=track, + save_path=save_path, save_dir=save_dir, stop_event=self.DL_POOL_STOP, progress=progress, @@ -745,6 +746,7 @@ class dl: elif track.descriptor == track.Descriptor.MPD: DASH.download_track( track=track, + save_path=save_path, save_dir=save_dir, stop_event=self.DL_POOL_STOP, progress=progress, @@ -791,13 +793,6 @@ class dl: track.drm = None if callable(track.OnDecrypted): track.OnDecrypted(track) - else: - with open(save_path, "wb") as f: - for file in sorted(save_dir.iterdir()): - f.write(file.read_bytes()) - file.unlink() - save_dir.rmdir() - track.path = save_path if track.path.stat().st_size <= 3: # Empty UTF-8 BOM == 3 bytes raise IOError( diff --git a/devine/core/manifests/dash.py b/devine/core/manifests/dash.py index 5295d97..7a104a8 100644 --- a/devine/core/manifests/dash.py +++ b/devine/core/manifests/dash.py @@ -4,6 +4,7 @@ import base64 import logging import math import re +import shutil import sys import time import traceback @@ -272,6 +273,7 @@ class DASH: @staticmethod def download_track( track: AnyTrack, + save_path: Path, save_dir: Path, stop_event: Event, progress: partial, @@ -504,13 +506,13 @@ class DASH: progress(total=len(segments)) + finished_threads = 0 + has_stopped = False + has_failed = False download_sizes = [] last_speed_refresh = time.time() with ThreadPoolExecutor(max_workers=16) as pool: - finished_threads = 0 - has_stopped = False - has_failed = False for download in futures.as_completed(( pool.submit( download_segment, @@ -556,10 +558,23 @@ class DASH: progress(downloaded=f"DASH {filesize.decimal(download_speed)}/s") last_speed_refresh = now download_sizes.clear() - if has_failed: - progress(downloaded="[red]FAILED") + + try: if has_stopped: progress(downloaded="[yellow]STOPPED") + return + if has_failed: + progress(downloaded="[red]FAILED") + return + + with open(save_path, "wb") as f: + for segment_file in sorted(save_dir.iterdir()): + f.write(segment_file.read_bytes()) + segment_file.unlink() + + track.path = save_path + finally: + shutil.rmtree(save_dir) @staticmethod def get_language(*options: Any) -> Optional[Language]: diff --git a/devine/core/manifests/hls.py b/devine/core/manifests/hls.py index 6d8eeaf..ed99f7d 100644 --- a/devine/core/manifests/hls.py +++ b/devine/core/manifests/hls.py @@ -2,6 +2,7 @@ from __future__ import annotations import logging import re +import shutil import sys import time import traceback @@ -181,6 +182,7 @@ class HLS: @staticmethod def download_track( track: AnyTrack, + save_path: Path, save_dir: Path, stop_event: Event, progress: partial, @@ -342,13 +344,13 @@ class HLS: progress(total=len(master.segments)) + finished_threads = 0 + has_stopped = False + has_failed = False download_sizes = [] last_speed_refresh = time.time() with ThreadPoolExecutor(max_workers=16) as pool: - finished_threads = 0 - has_stopped = False - has_failed = False for download in futures.as_completed(( pool.submit( download_segment, @@ -396,10 +398,23 @@ class HLS: progress(downloaded=f"HLS {filesize.decimal(download_speed)}/s") last_speed_refresh = now download_sizes.clear() - if has_failed: - progress(downloaded="[red]FAILED") + + try: if has_stopped: progress(downloaded="[yellow]STOPPED") + return + if has_failed: + progress(downloaded="[red]FAILED") + return + + with open(save_path, "wb") as f: + for segment_file in sorted(save_dir.iterdir()): + f.write(segment_file.read_bytes()) + segment_file.unlink() + + track.path = save_path + finally: + shutil.rmtree(save_dir) @staticmethod def get_drm(