Move segment merging from dl to DASH/HLS classes

This commit is contained in:
rlaphoenix 2023-03-01 08:54:35 +00:00
parent d07fedbbe1
commit 840db6e689
3 changed files with 42 additions and 17 deletions

View File

@ -735,6 +735,7 @@ class dl:
if track.descriptor == track.Descriptor.M3U: if track.descriptor == track.Descriptor.M3U:
HLS.download_track( HLS.download_track(
track=track, track=track,
save_path=save_path,
save_dir=save_dir, save_dir=save_dir,
stop_event=self.DL_POOL_STOP, stop_event=self.DL_POOL_STOP,
progress=progress, progress=progress,
@ -745,6 +746,7 @@ class dl:
elif track.descriptor == track.Descriptor.MPD: elif track.descriptor == track.Descriptor.MPD:
DASH.download_track( DASH.download_track(
track=track, track=track,
save_path=save_path,
save_dir=save_dir, save_dir=save_dir,
stop_event=self.DL_POOL_STOP, stop_event=self.DL_POOL_STOP,
progress=progress, progress=progress,
@ -791,13 +793,6 @@ class dl:
track.drm = None track.drm = None
if callable(track.OnDecrypted): if callable(track.OnDecrypted):
track.OnDecrypted(track) 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 if track.path.stat().st_size <= 3: # Empty UTF-8 BOM == 3 bytes
raise IOError( raise IOError(

View File

@ -4,6 +4,7 @@ import base64
import logging import logging
import math import math
import re import re
import shutil
import sys import sys
import time import time
import traceback import traceback
@ -272,6 +273,7 @@ class DASH:
@staticmethod @staticmethod
def download_track( def download_track(
track: AnyTrack, track: AnyTrack,
save_path: Path,
save_dir: Path, save_dir: Path,
stop_event: Event, stop_event: Event,
progress: partial, progress: partial,
@ -504,13 +506,13 @@ class DASH:
progress(total=len(segments)) progress(total=len(segments))
finished_threads = 0
has_stopped = False
has_failed = False
download_sizes = [] download_sizes = []
last_speed_refresh = time.time() last_speed_refresh = time.time()
with ThreadPoolExecutor(max_workers=16) as pool: with ThreadPoolExecutor(max_workers=16) as pool:
finished_threads = 0
has_stopped = False
has_failed = False
for download in futures.as_completed(( for download in futures.as_completed((
pool.submit( pool.submit(
download_segment, download_segment,
@ -556,10 +558,23 @@ class DASH:
progress(downloaded=f"DASH {filesize.decimal(download_speed)}/s") progress(downloaded=f"DASH {filesize.decimal(download_speed)}/s")
last_speed_refresh = now last_speed_refresh = now
download_sizes.clear() download_sizes.clear()
if has_failed:
progress(downloaded="[red]FAILED") try:
if has_stopped: if has_stopped:
progress(downloaded="[yellow]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 @staticmethod
def get_language(*options: Any) -> Optional[Language]: def get_language(*options: Any) -> Optional[Language]:

View File

@ -2,6 +2,7 @@ from __future__ import annotations
import logging import logging
import re import re
import shutil
import sys import sys
import time import time
import traceback import traceback
@ -181,6 +182,7 @@ class HLS:
@staticmethod @staticmethod
def download_track( def download_track(
track: AnyTrack, track: AnyTrack,
save_path: Path,
save_dir: Path, save_dir: Path,
stop_event: Event, stop_event: Event,
progress: partial, progress: partial,
@ -342,13 +344,13 @@ class HLS:
progress(total=len(master.segments)) progress(total=len(master.segments))
finished_threads = 0
has_stopped = False
has_failed = False
download_sizes = [] download_sizes = []
last_speed_refresh = time.time() last_speed_refresh = time.time()
with ThreadPoolExecutor(max_workers=16) as pool: with ThreadPoolExecutor(max_workers=16) as pool:
finished_threads = 0
has_stopped = False
has_failed = False
for download in futures.as_completed(( for download in futures.as_completed((
pool.submit( pool.submit(
download_segment, download_segment,
@ -396,10 +398,23 @@ class HLS:
progress(downloaded=f"HLS {filesize.decimal(download_speed)}/s") progress(downloaded=f"HLS {filesize.decimal(download_speed)}/s")
last_speed_refresh = now last_speed_refresh = now
download_sizes.clear() download_sizes.clear()
if has_failed:
progress(downloaded="[red]FAILED") try:
if has_stopped: if has_stopped:
progress(downloaded="[yellow]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 @staticmethod
def get_drm( def get_drm(