From f85ddce6f24be6fea490c9661c037af5d929f452 Mon Sep 17 00:00:00 2001 From: kenzuyaa Date: Tue, 26 Aug 2025 17:58:23 +0700 Subject: [PATCH] feat(downloaders): improve aria2c download progress reporting - Added RPC calls to get detailed global and active download statistics - Calculated total downloaded size, content size, and active download speed from active downloads - Included stopped downloads in totals and handle error states with logged messages - Yielded enhanced progress updates with combined downloaded sizes and speeds - Added more granular progress dictionary keys for richer status reporting - Added sleep delay in main aria2c function to reduce CPU usage during monitoring loop - Updated docstring examples to reflect new progress data format and keys --- unshackle/core/downloaders/aria2c.py | 70 ++++++++++++++++++++++------ 1 file changed, 57 insertions(+), 13 deletions(-) diff --git a/unshackle/core/downloaders/aria2c.py b/unshackle/core/downloaders/aria2c.py index bc43460..3926d58 100644 --- a/unshackle/core/downloaders/aria2c.py +++ b/unshackle/core/downloaders/aria2c.py @@ -137,7 +137,6 @@ def download( if len(urls) > 1: split = 1 file_allocation = "none" - arguments = [ # [Basic Options] "--input-file", @@ -189,19 +188,36 @@ def download( p.stdin.close() while p.poll() is None: + # Get global statistics via RPC global_stats: dict[str, Any] = ( rpc(caller=partial(rpc_session.post, url=rpc_uri), secret=rpc_secret, method="aria2.getGlobalStat") or {} ) number_stopped = int(global_stats.get("numStoppedTotal", 0)) - download_speed = int(global_stats.get("downloadSpeed", -1)) + global_download_speed = int(global_stats.get("downloadSpeed", 0)) - if number_stopped: - yield dict(completed=number_stopped) - if download_speed != -1: - yield dict(downloaded=f"{filesize.decimal(download_speed)}/s") + # Get active downloads via RPC for detailed progress tracking + active_downloads: list[dict[str, Any]] = ( + rpc( + caller=partial(rpc_session.post, url=rpc_uri), + secret=rpc_secret, + method="aria2.tellActive", + ) + or [] + ) + # Calculate totals from active downloads + total_downloaded_size = 0 + total_content_size = 0 + active_download_speed = 0 + + for download in active_downloads: + total_downloaded_size += int(download.get("completedLength", 0)) + total_content_size += int(download.get("totalLength", 0)) + active_download_speed += int(download.get("downloadSpeed", 0)) + + # Get stopped downloads via RPC to check for errors and completion stopped_downloads: list[dict[str, Any]] = ( rpc( caller=partial(rpc_session.post, url=rpc_uri), @@ -212,22 +228,46 @@ def download( or [] ) - for dl in stopped_downloads: - if dl["status"] == "error": + # Add completed downloads to totals and handle errors + for download in stopped_downloads: + if download["status"] == "complete": + completed_length = int(download.get("completedLength", 0)) + total_downloaded_size += completed_length + total_content_size += completed_length + elif download["status"] == "error": used_uri = next( uri["uri"] - for file in dl["files"] + for file in download["files"] if file["selected"] == "true" for uri in file["uris"] if uri["status"] == "used" ) - error = f"Download Error (#{dl['gid']}): {dl['errorMessage']} ({dl['errorCode']}), {used_uri}" + error = f"Download Error (#{download['gid']}): {download['errorMessage']} ({download['errorCode']}), {used_uri}" error_pretty = "\n ".join( textwrap.wrap(error, width=console.width - 20, initial_indent="") ) console.log(Text.from_ansi("\n[Aria2c]: " + error_pretty)) raise ValueError(error) + # Yield progress information + if total_content_size > 0: + downloaded = f"{filesize.decimal(total_downloaded_size)}/{filesize.decimal(total_content_size)} {filesize.decimal(active_download_speed)}/s" + yield dict( + downloaded=downloaded, + total=total_content_size, + completed=total_downloaded_size + ) + elif global_download_speed > 0: + yield dict( + downloaded=f"{filesize.decimal(global_download_speed)}/s", + speed_bytes_per_sec=global_download_speed + ) + + # Yield completion count + if number_stopped: + yield dict(completed=number_stopped) + + # Exit when all downloads are complete if number_stopped == len(urls): rpc(caller=partial(rpc_session.post, url=rpc_uri), secret=rpc_secret, method="aria2.shutdown") break @@ -274,11 +314,14 @@ def aria2c( Yields the following download status updates while chunks are downloading: - - {total: 100} (100% download total) - - {completed: 1} (1% download progress out of 100%) - - {downloaded: "10.1 MB/s"} (currently downloading at a rate of 10.1 MB/s) + - {total: 100} (total number of URLs to download) + - {completed: 1} (number of completed downloads) + - {downloaded: "50.2 MB/128.5 MB 10.1 MB/s", total: 134742016, completed: 52428800} (progress data) + - {downloaded: "10.1 MB/s", speed_bytes_per_sec: 10485760} (speed fallback data) The data is in the same format accepted by rich's progress.update() function. + However, The `downloaded` and `speed_bytes_per_sec` keys are custom and not natively + accepted by all rich progress bars. Parameters: urls: Web URL(s) to file(s) to download. You can use a dictionary with the key @@ -318,6 +361,7 @@ def aria2c( stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) + time.sleep(1) try: yield from download(urls, output_dir, filename, headers, cookies, local_proxy, max_workers)