Fix printing of aria2c errors if progress is used

Also improves the general process of ingesting aria2c progress information.
This commit is contained in:
rlaphoenix 2023-02-27 16:52:00 +00:00
parent 18449c4777
commit e5e3f4687d
1 changed files with 39 additions and 33 deletions

View File

@ -1,12 +1,14 @@
import asyncio import asyncio
import subprocess import subprocess
import sys
from asyncio import IncompleteReadError from asyncio import IncompleteReadError
from functools import partial from functools import partial
from pathlib import Path from pathlib import Path
from typing import Optional, Union from typing import Optional, Union
from rich.text import Text
from devine.core.config import config from devine.core.config import config
from devine.core.console import console
from devine.core.utilities import get_binary_path, start_pproxy from devine.core.utilities import get_binary_path, start_pproxy
@ -100,41 +102,45 @@ async def aria2c(
p.stdin.close() p.stdin.close()
if progress: if progress:
def update_progress_bar(data: str):
if "%" in data:
# id, dledMiB/totalMiB(x%), CN:xx, DL:xxMiB, ETA:Xs
# eta may not always be available
data_parts = data[1:-1].split()
perc_parts = data_parts[1].split("(")
if len(perc_parts) == 2:
# might otherwise be e.g., 0B/0B, with no % symbol provided
progress(
total=100,
completed=int(perc_parts[1][:-2]),
downloaded=f"{data_parts[3].split(':')[1]}/s"
)
# I'm sorry for this shameful code, aria2(c) is annoying as f!!! # I'm sorry for this shameful code, aria2(c) is annoying as f!!!
buffer = b""
recording = False
while not p.stdout.at_eof(): while not p.stdout.at_eof():
try: try:
byte = await p.stdout.readexactly(1) buffer = await p.stdout.readuntil(b"\r")
except IncompleteReadError: except IncompleteReadError as e:
pass # ignore, the first read will do this buffer = e.partial
else:
if byte == b"=": # download result log buffer = buffer.decode().strip()
progress(total=100, completed=100) if buffer:
break buffer_lines = buffer.splitlines()
if byte == b"[": is_dl_summary = False
recording = True for line in buffer_lines:
if recording: if line:
buffer += byte if line.startswith("[") and line.endswith("]"):
if byte == b"]": update_progress_bar(line)
recording = False elif line.startswith("Download Results"):
if b"FileAlloc" not in buffer and b"ERROR" not in buffer: # we know it's 100% downloaded, but let's use the avg dl speed value
try: is_dl_summary = True
# id, dledMiB/totalMiB(x%), CN:xx, DL:xxMiB, ETA:Xs elif is_dl_summary and "OK" in line and "|" in line:
# eta may not always be available gid, status, avg_speed, path_or_uri = line.split("|")
parts = buffer.decode()[1:-1].split() progress(total=100, completed=100, downloaded=avg_speed.strip())
dl_parts = parts[1].split("(") elif not is_dl_summary:
if len(dl_parts) == 2: buffer_msg = line.split(" ", maxsplit=2)
# might otherwise be e.g., 0B/0B, with no % symbol provided buffer_msg = f"[Aria2c]: {buffer_msg[-1].strip()}"
progress( console.log(Text.from_ansi(buffer_msg))
total=100,
completed=int(dl_parts[1][:-2]),
downloaded=f"{parts[3].split(':')[1]}/s"
)
except Exception as e:
print(f"Aria2c progress failed on {buffer}, {e!r}")
sys.exit(1)
buffer = b""
await p.wait() await p.wait()