Add ability to set extra arguments with aria2c

This commit is contained in:
rlaphoenix 2023-02-22 03:01:54 +00:00
parent 1443dfaecc
commit 0913b0dda6
1 changed files with 19 additions and 12 deletions

View File

@ -14,7 +14,8 @@ async def aria2c(
out: Path, out: Path,
headers: Optional[dict] = None, headers: Optional[dict] = None,
proxy: Optional[str] = None, proxy: Optional[str] = None,
byte_range: Optional[str] = None byte_range: Optional[str] = None,
*args: str
) -> int: ) -> int:
""" """
Download files using Aria2(c). Download files using Aria2(c).
@ -61,30 +62,36 @@ async def aria2c(
"--file-allocation", config.aria2c.get("file_allocation", "falloc"), "--file-allocation", config.aria2c.get("file_allocation", "falloc"),
"--console-log-level", "warn", "--console-log-level", "warn",
"--download-result", "hide", "--download-result", "hide",
*args,
"-i", "-" "-i", "-"
] ]
for header, value in (headers or {}).items(): headers = headers or {}
if byte_range:
headers["Range"] = f"bytes={byte_range}"
for header, value in headers.items():
if header.lower() == "accept-encoding": if header.lower() == "accept-encoding":
# we cannot set an allowed encoding, or it will return compressed # we cannot set an allowed encoding, or it will return compressed
# and the code is not set up to uncompress the data # and the code is not set up to uncompress the data
continue continue
arguments.extend(["--header", f"{header}: {value}"]) arguments.extend(["--header", f"{header}: {value}"])
if byte_range:
arguments.extend(["--header", f"Range: bytes={byte_range}"])
if proxy and proxy.lower().split(":")[0] != "http":
# HTTPS proxies not supported by Aria2c.
# Proxy the proxy via pproxy to access it as a HTTP proxy.
async with start_pproxy(proxy) as pproxy_:
return await aria2c(uri, out, headers, pproxy_)
if proxy: 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.
async with start_pproxy(proxy) as pproxy_:
return await aria2c(uri, out, headers, pproxy_)
arguments += ["--all-proxy", proxy] arguments += ["--all-proxy", proxy]
p = await asyncio.create_subprocess_exec(executable, *arguments, stdin=subprocess.PIPE) p = await asyncio.create_subprocess_exec(
executable,
*arguments,
stdin=subprocess.PIPE
)
await p.communicate(uri.encode()) await p.communicate(uri.encode())
if p.returncode != 0: if p.returncode != 0:
raise subprocess.CalledProcessError(p.returncode, arguments) raise subprocess.CalledProcessError(p.returncode, arguments)