150 lines
5.6 KiB
Python
150 lines
5.6 KiB
Python
import subprocess
|
|
import textwrap
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import Optional, Union
|
|
|
|
from rich.text import Text
|
|
|
|
from devine.core.config import config
|
|
from devine.core.console import console
|
|
from devine.core.utilities import get_binary_path, start_pproxy
|
|
|
|
|
|
def aria2c(
|
|
uri: Union[str, list[str]],
|
|
out: Path,
|
|
headers: Optional[dict] = None,
|
|
proxy: Optional[str] = None,
|
|
silent: bool = False,
|
|
progress: Optional[partial] = None,
|
|
*args: str
|
|
) -> int:
|
|
"""
|
|
Download files using Aria2(c).
|
|
https://aria2.github.io
|
|
|
|
If multiple URLs are provided they will be downloaded in the provided order
|
|
to the output directory. They will not be merged together.
|
|
"""
|
|
segmented = False
|
|
if isinstance(uri, list) and len(uri) == 1:
|
|
uri = uri[0]
|
|
if isinstance(uri, list):
|
|
segmented = True
|
|
uri = "\n".join([
|
|
f"{url}\n"
|
|
f"\tdir={out}\n"
|
|
f"\tout={i:08}.mp4"
|
|
for i, url in enumerate(uri)
|
|
])
|
|
if out.is_file():
|
|
raise ValueError("Provided multiple segments to download, expecting directory path")
|
|
elif "\t" not in uri:
|
|
uri = f"{uri}\n" \
|
|
f"\tdir={out.parent}\n" \
|
|
f"\tout={out.name}"
|
|
|
|
executable = get_binary_path("aria2c", "aria2")
|
|
if not executable:
|
|
raise EnvironmentError("Aria2c executable not found...")
|
|
|
|
arguments = [
|
|
"-c", # Continue downloading a partially downloaded file
|
|
"--remote-time", # Retrieve timestamp of the remote file from the and apply if available
|
|
"-x", "16", # The maximum number of connections to one server for each download
|
|
"-j", "16", # The maximum number of parallel downloads for every static (HTTP/FTP) URL
|
|
"-s", ("1" if segmented else "16"), # Download a file using N connections
|
|
"--min-split-size", ("1024M" if segmented else "20M"), # effectively disable split if segmented
|
|
"--allow-overwrite=true",
|
|
"--auto-file-renaming=false",
|
|
"--retry-wait", "2", # Set the seconds to wait between retries.
|
|
"--max-tries", "5",
|
|
"--max-file-not-found", "5",
|
|
"--summary-interval", "0",
|
|
"--file-allocation", config.aria2c.get("file_allocation", "prealloc"),
|
|
"--console-log-level", "warn",
|
|
"--download-result", ["hide", "default"][bool(progress)],
|
|
*args,
|
|
"-i", "-"
|
|
]
|
|
|
|
for header, value in (headers or {}).items():
|
|
if header.lower() == "accept-encoding":
|
|
# we cannot set an allowed encoding, or it will return compressed
|
|
# and the code is not set up to uncompress the data
|
|
continue
|
|
arguments.extend(["--header", f"{header}: {value}"])
|
|
|
|
if proxy:
|
|
if proxy.lower().split(":")[0] != "http":
|
|
# HTTPS proxies are not supported by aria2(c).
|
|
# Proxy the proxy via pproxy to access it as an HTTP proxy.
|
|
with start_pproxy(proxy) as pproxy_:
|
|
return aria2c(uri, out, headers, pproxy_)
|
|
arguments += ["--all-proxy", proxy]
|
|
|
|
try:
|
|
p = subprocess.Popen(
|
|
[executable, *arguments],
|
|
stdin=subprocess.PIPE,
|
|
stdout=[subprocess.PIPE, subprocess.DEVNULL][silent],
|
|
universal_newlines=True
|
|
)
|
|
p._stdin_write(uri) # noqa
|
|
|
|
is_dl_summary = False
|
|
aria_log_buffer = ""
|
|
for line in iter(p.stdout.readline, ""):
|
|
line = line.strip()
|
|
if line:
|
|
if line.startswith("Download Results"):
|
|
# we know it's 100% downloaded, but let's use the avg dl speed value
|
|
is_dl_summary = True
|
|
elif line.startswith("[") and line.endswith("]"):
|
|
if progress and "%" in line:
|
|
# id, dledMiB/totalMiB(x%), CN:xx, DL:xxMiB, ETA:Xs
|
|
# eta may not always be available
|
|
data_parts = line[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"
|
|
)
|
|
elif is_dl_summary and "OK" in line and "|" in line:
|
|
gid, status, avg_speed, path_or_uri = line.split("|")
|
|
progress(total=100, completed=100, downloaded=avg_speed.strip())
|
|
elif not is_dl_summary:
|
|
aria_log_buffer += f"{line.strip()}\n"
|
|
|
|
if aria_log_buffer:
|
|
# wrap to console width - padding - '[Aria2c]: '
|
|
aria_log_buffer = "\n ".join(textwrap.wrap(
|
|
aria_log_buffer.rstrip(),
|
|
width=console.width - 20,
|
|
initial_indent=""
|
|
))
|
|
console.log(Text.from_ansi("\n[Aria2c]: " + aria_log_buffer))
|
|
|
|
p.wait()
|
|
|
|
if p.returncode != 0:
|
|
raise subprocess.CalledProcessError(p.returncode, arguments)
|
|
except ConnectionResetError:
|
|
# interrupted while passing URI to download
|
|
raise KeyboardInterrupt()
|
|
except subprocess.CalledProcessError as e:
|
|
if e.returncode in (7, 0xC000013A):
|
|
# 7 is when Aria2(c) handled the CTRL+C
|
|
# 0xC000013A is when it never got the chance to
|
|
raise KeyboardInterrupt()
|
|
raise
|
|
|
|
return p.returncode
|
|
|
|
|
|
__ALL__ = (aria2c,)
|