ninjajiraiya
/
devine
forked from DRMTalks/devine
1
0
Fork 0
Code Pull Requests Activity

feat(dl): Support multiple -r/--range and mux ranges separately 

Multiple -r/--range values can be used with multiple -q/--quality values.

Closes #63
This commit is contained in:
rlaphoenix 2024-03-04 13:11:43 +00:00
parent 6e8efc3f63
commit 0201c41feb
1 changed files with 56 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

87
devine/commands/dl.py
View File

@ -14,7 +14,7 @@ from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy from copy import deepcopy
from functools import partial from functools import partial
from http.cookiejar import CookieJar, MozillaCookieJar from http.cookiejar import CookieJar, MozillaCookieJar
from itertools import zip_longest from itertools import product
from pathlib import Path from pathlib import Path
from threading import Lock from threading import Lock
from typing import Any, Callable, Optional from typing import Any, Callable, Optional
@ -32,7 +32,7 @@ from rich.console import Group
from rich.live import Live from rich.live import Live
from rich.padding import Padding from rich.padding import Padding
from rich.panel import Panel from rich.panel import Panel
from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn, TimeRemainingColumn from rich.progress import BarColumn, Progress, SpinnerColumn, TaskID, TextColumn, TimeRemainingColumn
from rich.rule import Rule from rich.rule import Rule
from rich.table import Table from rich.table import Table
from rich.text import Text from rich.text import Text
@ -50,7 +50,7 @@ from devine.core.titles import Movie, Song, Title_T
from devine.core.titles.episode import Episode from devine.core.titles.episode import Episode
from devine.core.tracks import Audio, Subtitle, Tracks, Video from devine.core.tracks import Audio, Subtitle, Tracks, Video
from devine.core.utilities import get_binary_path, is_close_match, time_elapsed_since 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.click_types import LANGUAGE_RANGE, QUALITY_LIST, SEASON_RANGE, ContextData, MultipleChoice
from devine.core.utils.collections import merge_dict from devine.core.utils.collections import merge_dict
from devine.core.utils.subprocess import ffprobe from devine.core.utils.subprocess import ffprobe
from devine.core.vaults import Vaults from devine.core.vaults import Vaults
@ -81,9 +81,9 @@ class dl:
@click.option("-ab", "--abitrate", type=int, @click.option("-ab", "--abitrate", type=int,
default=None, default=None,
help="Audio Bitrate to download (in kbps), defaults to highest available.") help="Audio Bitrate to download (in kbps), defaults to highest available.")
@click.option("-r", "--range", "range_", type=click.Choice(Video.Range, case_sensitive=False), @click.option("-r", "--range", "range_", type=MultipleChoice(Video.Range, case_sensitive=False),
default=Video.Range.SDR, default=[Video.Range.SDR],
help="Video Color Range, defaults to SDR.") help="Video Color Range(s) to download, defaults to SDR.")
@click.option("-c", "--channels", type=float, @click.option("-c", "--channels", type=float,
default=None, default=None,
help="Audio Channel(s) to download. Matches sub-channel layouts like 5.1 with 6.0 implicitly.") help="Audio Channel(s) to download. Matches sub-channel layouts like 5.1 with 6.0 implicitly.")
@ -254,7 +254,7 @@ class dl:
acodec: Optional[Audio.Codec], acodec: Optional[Audio.Codec],
vbitrate: int, vbitrate: int,
abitrate: int, abitrate: int,
range_: Video.Range, range_: list[Video.Range],
channels: float, channels: float,
wanted: list[str], wanted: list[str],
lang: list[str], lang: list[str],
@ -363,10 +363,12 @@ class dl:
self.log.error(f"There's no {vcodec.name} Video Track...") self.log.error(f"There's no {vcodec.name} Video Track...")
sys.exit(1) sys.exit(1)
title.tracks.select_video(lambda x: x.range == range_) if range_:
if not title.tracks.videos: title.tracks.select_video(lambda x: x.range in range_)
self.log.error(f"There's no {range_.name} Video Track...") for color_range in range_:
sys.exit(1) if not any(x.range == color_range for x in title.tracks.videos):
self.log.error(f"There's no {color_range.name} Video Tracks...")
sys.exit(1)
if vbitrate: if vbitrate:
title.tracks.select_video(lambda x: x.bitrate and x.bitrate // 1000 == vbitrate) title.tracks.select_video(lambda x: x.bitrate and x.bitrate // 1000 == vbitrate)
@ -382,7 +384,7 @@ class dl:
sys.exit(1) sys.exit(1)
if quality: if quality:
title.tracks.by_resolutions(quality, per_resolution=1) title.tracks.by_resolutions(quality)
missing_resolutions = [] missing_resolutions = []
for resolution in quality: for resolution in quality:
if any(video.height == resolution for video in title.tracks.videos): if any(video.height == resolution for video in title.tracks.videos):
@ -398,8 +400,27 @@ class dl:
plural = "s" if len(missing_resolutions) > 1 else "" plural = "s" if len(missing_resolutions) > 1 else ""
self.log.error(f"There's no {res_list} Video Track{plural}...") self.log.error(f"There's no {res_list} Video Track{plural}...")
sys.exit(1) sys.exit(1)
else:
title.tracks.videos = [title.tracks.videos[0]] # choose best track by range and quality
title.tracks.videos = [
track
for resolution, color_range in product(
quality or [None],
range_ or [None]
)
for track in [next(
t
for t in title.tracks.videos
if (not resolution and not color_range) or
(
(not resolution or (
(t.height == resolution) or
(int(t.width * (9 / 16)) == resolution)
))
and (not color_range or t.range == color_range)
)
)]
]
# filter subtitle tracks # filter subtitle tracks
if s_lang and "all" not in s_lang: if s_lang and "all" not in s_lang:
@ -599,26 +620,31 @@ class dl:
TimeRemainingColumn(compact=True, elapsed_when_finished=True), TimeRemainingColumn(compact=True, elapsed_when_finished=True),
console=console console=console
) )
multi_jobs = len(title.tracks.videos) > 1
tasks = [ multiplex_tasks: list[tuple[TaskID, Tracks]] = []
progress.add_task( for video_track in title.tracks.videos:
f"Multiplexing{f' {x.height}p' if multi_jobs else ''}...", task_description = "Multiplexing"
total=None, if len(quality) > 1:
start=False task_description += f" {video_track.height}p"
) if len(range_) > 1:
for x in title.tracks.videos or [None] task_description += f" {video_track.range.name}"
]
task_id = progress.add_task(f"{task_description}...", total=None, start=False)
task_tracks = Tracks(title.tracks)
task_tracks.videos = [video_track]
multiplex_tasks.append((task_id, task_tracks))
with Live( with Live(
Padding(progress, (0, 5, 1, 5)), Padding(progress, (0, 5, 1, 5)),
console=console console=console
): ):
for task, video_track in zip_longest(tasks, title.tracks.videos, fillvalue=None): for task_id, task_tracks in multiplex_tasks:
if video_track: progress.start_task(task_id) # TODO: Needed?
title.tracks.videos = [video_track] muxed_path, return_code = task_tracks.mux(
progress.start_task(task) # TODO: Needed?
muxed_path, return_code = title.tracks.mux(
str(title), str(title),
progress=partial(progress.update, task_id=task), progress=partial(progress.update, task_id=task_id),
delete=False delete=False
) )
muxed_paths.append(muxed_path) muxed_paths.append(muxed_path)
@ -627,8 +653,7 @@ class dl:
elif return_code >= 2: elif return_code >= 2:
self.log.error(f"Failed to Mux video to Matroska file ({return_code})") self.log.error(f"Failed to Mux video to Matroska file ({return_code})")
sys.exit(1) sys.exit(1)
if video_track: task_tracks.videos[0].delete()
video_track.delete()
for track in title.tracks: for track in title.tracks:
track.delete() track.delete()
else: else: