Add rich progress bar for the mux operation

Update dl.py
This commit is contained in:
rlaphoenix 2023-02-25 13:58:58 +00:00
parent 58673590df
commit a5da7c8fbd
2 changed files with 43 additions and 17 deletions

View File

@ -29,6 +29,7 @@ from pywidevine.device import Device
from pywidevine.remotecdm import RemoteCdm from pywidevine.remotecdm import RemoteCdm
from rich.live import Live from rich.live import Live
from rich.padding import Padding from rich.padding import Padding
from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn, TimeRemainingColumn
from rich.rule import Rule from rich.rule import Rule
from devine.core.config import config from devine.core.config import config
@ -733,16 +734,32 @@ class dl:
def mux_tracks(self, title: Title_T, season_folder: bool = True, add_source: bool = True) -> None: def mux_tracks(self, title: Title_T, season_folder: bool = True, add_source: bool = True) -> None:
"""Mux Tracks, Delete Pre-Mux files, and move to the final location.""" """Mux Tracks, Delete Pre-Mux files, and move to the final location."""
console.log("Muxing Tracks into a Matroska Container")
if isinstance(title, (Movie, Episode)): if isinstance(title, (Movie, Episode)):
muxed_path, return_code = title.tracks.mux(str(title)) multiplexing_progress = Progress(
if return_code == 1: TextColumn("[progress.description]{task.description}"),
self.log.warning("mkvmerge had at least one warning, will continue anyway...") SpinnerColumn(),
elif return_code >= 2: BarColumn(),
self.log.error(" - Failed to Mux video to Matroska file") "",
sys.exit(1) TimeRemainingColumn(compact=True, elapsed_when_finished=True),
console.log(f" + Muxed to {muxed_path}") console=console
)
with Live(
Padding(multiplexing_progress, (0, 5, 1, 5)),
console=console
):
task = multiplexing_progress.add_task("Multiplexing...", total=100)
muxed_path, return_code = title.tracks.mux(
str(title),
progress=partial(
multiplexing_progress.update,
task_id=task
)
)
if return_code == 1:
self.log.warning("mkvmerge had at least one warning, will continue anyway...")
elif return_code >= 2:
self.log.error(" - Failed to Mux video to Matroska file")
sys.exit(1)
else: else:
# dont mux # dont mux
muxed_path = title.tracks.audio[0].path muxed_path = title.tracks.audio[0].path
@ -758,7 +775,6 @@ class dl:
final_path = final_dir / f"{final_filename}{muxed_path.suffix}" final_path = final_dir / f"{final_filename}{muxed_path.suffix}"
shutil.move(muxed_path, final_path) shutil.move(muxed_path, final_path)
console.log(f" + Moved to {final_path}")
@staticmethod @staticmethod
def get_profile(service: str) -> Optional[str]: def get_profile(service: str) -> Optional[str]:

View File

@ -292,10 +292,16 @@ class Tracks:
tracks_.append(next(x for x in tracks if str(x.language) == match)) tracks_.append(next(x for x in tracks if str(x.language) == match))
return tracks_ return tracks_
def mux(self, title: str, delete: bool = True) -> tuple[Path, int]: def mux(self, title: str, delete: bool = True, progress: Optional[partial] = None) -> tuple[Path, int]:
""" """
Takes the Video, Audio and Subtitle Tracks, and muxes them into an MKV file. Multiplex all the Tracks into a Matroska Container file.
It will attempt to detect Forced/Default tracks, and will try to parse the language codes of the Tracks
Parameters:
title: Set the Matroska Container file title. Usually displayed in players
instead of the filename if set.
delete: Delete all track files after multiplexing.
progress: Update a rich progress bar via `completed=...`. This must be the
progress object's update() func, pre-set with task id via functools.partial.
""" """
cl = [ cl = [
"mkvmerge", "mkvmerge",
@ -373,11 +379,15 @@ class Tracks:
# let potential failures go to caller, caller should handle # let potential failures go to caller, caller should handle
try: try:
p = subprocess.run([ p = subprocess.Popen([
*cl, *cl,
"--output", str(output_path) "--output", str(output_path),
]) "--gui-mode"
return output_path, p.returncode ], text=True, stdout=subprocess.PIPE)
for line in iter(p.stdout.readline, ""):
if "progress" in line:
progress(total=100, completed=int(line.strip()[14:-1]))
return output_path, p.wait()
finally: finally:
if chapters_path: if chapters_path:
# regardless of delete param, we delete as it's a file we made during muxing # regardless of delete param, we delete as it's a file we made during muxing