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:
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(

View File

@ -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]:

View File

@ -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(