Clean up residual files on download stops and fails

This commit is contained in:
rlaphoenix 2023-03-01 11:17:46 +00:00
parent 9f48aab80c
commit 7b7be47f7d
1 changed files with 74 additions and 66 deletions

View File

@ -8,7 +8,6 @@ import re
import shutil import shutil
import sys import sys
import time import time
import traceback
from concurrent import futures from concurrent import futures
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy from copy import deepcopy
@ -729,79 +728,88 @@ class dl:
else: else:
save_dir = save_path.parent save_dir = save_path.parent
def cleanup():
# e.g., foo.mp4
save_path.unlink(missing_ok=True)
# e.g., foo.mp4.aria2
save_path.with_suffix(f"{save_path.suffix}.aria2").unlink(missing_ok=True)
for file in config.directories.temp.glob(f"{save_path.stem}.*{save_path.suffix}"):
# e.g., foo.decrypted.mp4, foo.repack.mp4, and such
file.unlink()
if save_dir.exists() and save_dir.name.endswith("_segments"):
shutil.rmtree(save_dir)
# Delete any pre-existing temp files matching this track. # 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 # 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 # 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. # corruptions caused by sudden interruptions to the lock file.
for existing_file in config.directories.temp.glob(f"{save_path.stem}.*{save_path.suffix}"): cleanup()
# e.g., foo.decrypted.mp4, foo.repack.mp4, and such
existing_file.unlink()
if save_dir.exists() and save_dir.name.endswith("_segments"):
shutil.rmtree(save_dir)
if track.descriptor == track.Descriptor.M3U: try:
HLS.download_track( if track.descriptor == track.Descriptor.M3U:
track=track, HLS.download_track(
save_path=save_path, track=track,
save_dir=save_dir, save_path=save_path,
stop_event=self.DL_POOL_STOP, save_dir=save_dir,
progress=progress, stop_event=self.DL_POOL_STOP,
session=service.session, progress=progress,
proxy=proxy, session=service.session,
license_widevine=prepare_drm 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,
stop_event=self.DL_POOL_STOP,
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:
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?")
else:
prepare_drm(drm)
track.drm = [drm]
aria2c(
uri=track.url,
out=save_path,
headers=service.session.headers,
proxy=proxy if track.needs_proxy else None,
progress=progress
) )
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,
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:
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?")
else:
prepare_drm(drm)
track.drm = [drm]
track.path = save_path aria2c(
uri=track.url,
out=save_path,
headers=service.session.headers,
proxy=proxy if track.needs_proxy else None,
progress=progress
)
if track.drm: track.path = save_path
drm = track.drm[0] # just use the first supported DRM system for now
drm.decrypt(save_path) if track.drm:
track.drm = None drm = track.drm[0] # just use the first supported DRM system for now
if callable(track.OnDecrypted): drm.decrypt(save_path)
track.OnDecrypted(track) track.drm = None
except KeyboardInterrupt: if callable(track.OnDecrypted):
progress(downloaded="[yellow]STOPPED") track.OnDecrypted(track)
except Exception as e: except KeyboardInterrupt:
progress(downloaded="[red]FAILED") self.DL_POOL_STOP.set()
traceback.print_exception(e) progress(downloaded="[yellow]STOPPED")
self.log.error(f"URL Download worker threw an unhandled exception: {e!r}") raise
finally: except Exception:
self.DL_POOL_STOP.set() self.DL_POOL_STOP.set()
save_path.unlink(missing_ok=True) progress(downloaded="[red]FAILED")
save_path.with_suffix(f"{save_path.suffix}.aria2").unlink(missing_ok=True) raise
except (Exception, KeyboardInterrupt):
cleanup()
raise
if self.DL_POOL_STOP.is_set(): if self.DL_POOL_STOP.is_set():
# we stopped during the download, let's exit # we stopped during the download, let's exit