From 50ea658fc7fd7bbdcf53fa11068bcb71edcc759d Mon Sep 17 00:00:00 2001 From: Abdussamet Kocak Date: Tue, 20 Oct 2026 12:59:23 +0000 Subject: [PATCH] fix: Resume unfinished imageset downloads --- download_imageset.py | 45 ++++++++++++++++------ download_imageset_test.py | 80 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 12 deletions(-) mode change 100644 => 100755 download_imageset_test.py diff --git a/download_imageset.py b/download_imageset.py index 8516c96..1f75c51 100755 --- a/download_imageset.py +++ b/download_imageset.py @@ -284,6 +284,9 @@ def process_imageset(path: Path, processor: ImageProcessor) -> None: return for archive in archives: + if archive.with_name(f"{archive.name}.aria2").exists(): + logging.warning(f"skipping {archive}: download is unfinished") + return logging.info(f"extracting {archive} → {output_dir}") if archive.suffix.lower() == ".zip": try: @@ -342,7 +345,7 @@ def remove_empty_dirs(root_dir: Path) -> None: for d in removed: d.rmdir() if removed: - print(f"removed {len(removed)} empty dir(s) in {root_dir}") + logging.info(f"removed {len(removed)} empty dir(s) in {root_dir}") def move_to_trash(file_paths: list[Path]) -> None: @@ -423,9 +426,9 @@ def gallery_cmd(urls: tuple[str, ...], name: str | None, cwd: Path | None) -> No output_dir = (cwd or Path(".")) / resolved_name if output_dir.exists() and _has_images(output_dir): - print(f"skipping gallery-dl: {output_dir} already has images") + logging.info(f"skipping gallery-dl: {output_dir} already has images") else: - print(f"gallery-dl → {output_dir}") + logging.info(f"gallery-dl → {output_dir}") args = ["uvx", "gallery-dl"] if name or cwd: args += ["-D", str(output_dir)] @@ -445,24 +448,42 @@ def filehost_cmd(urls: tuple[str, ...], name: str | None, cwd: Path, token: str) output_dir.mkdir(parents=True, exist_ok=True) for url in urls: - print(f"unrestricting {url}") + logging.info(f"unrestricting {url}") dl_url = resolve_download_url_with_alldebrid(token, url) meta = _resolve_url(dl_url) if meta.suggested_extension in _VIDEO_EXTS: video_path = (cwd or Path(".")) / (resolved_name + meta.suggested_extension) - print(f"video detected → {video_path}") - if video_path.exists() and video_path.stat().st_size > 5 * 1024 * 1024: - print(f"skipping: {video_path.name} already downloaded") + logging.info(f"video detected → {video_path}") + if video_path.with_name(f"{video_path.name}.aria2").exists(): + logging.warning(f"unfinished download; resuming {video_path}") + if ( + video_path.exists() + and not video_path.with_name(f"{video_path.name}.aria2").exists() + and video_path.stat().st_size > 5 * 1024 * 1024 + ): + logging.info(f"skipping: {video_path.name} already downloaded") continue - run_command(["aria2c", "-o", str(video_path), dl_url]) + run_command(args=[ + "aria2c", "--continue=true", f"--dir={video_path.parent}", + f"--out={video_path.name}", dl_url, + ]) continue dest = output_dir / meta.suggested_filename - if dest.exists() and dest.stat().st_size > 5 * 1024 * 1024: - print(f"skipping: {dest.name} already downloaded") + if dest.with_name(f"{dest.name}.aria2").exists(): + logging.warning(f"unfinished download; resuming {dest}") + if ( + dest.exists() + and not dest.with_name(f"{dest.name}.aria2").exists() + and dest.stat().st_size > 5 * 1024 * 1024 + ): + logging.info(f"skipping: {dest.name} already downloaded") continue - print(f"downloading {meta.suggested_filename} → {output_dir}") - run_command(["aria2c", f"--dir={output_dir}", dl_url]) + logging.info(f"downloading {meta.suggested_filename} → {output_dir}") + run_command(args=[ + "aria2c", "--continue=true", f"--dir={output_dir}", + f"--out={dest.name}", dl_url, + ]) process_imageset(path=output_dir, processor=detect_resizer()) diff --git a/download_imageset_test.py b/download_imageset_test.py old mode 100644 new mode 100755 index 56121c4..58cace0 --- a/download_imageset_test.py +++ b/download_imageset_test.py @@ -297,5 +297,85 @@ def test_resume_resizing( trash.assert_called_once_with(file_paths=[image_path]) +@pytest.mark.parametrize("extension", [".zip", ".mp4"]) +@pytest.mark.parametrize("unfinished", [False, True]) +def test_filehost_resumes_unfinished_downloads( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + extension: str, + unfinished: bool, +) -> None: + directory = tmp_path / "album" + directory.mkdir() + destination = ( + directory / "download.zip" if extension == ".zip" else tmp_path / "album.mp4" + ) + with destination.open(mode="wb") as output: + output.truncate(6 * 1024 * 1024) + control_file = destination.with_name(f"{destination.name}.aria2") + if unfinished: + control_file.touch() + url = "https://example.com/download" + monkeypatch.setattr( + download_imageset, "resolve_download_url_with_alldebrid", Mock(return_value=url) + ) + monkeypatch.setattr( + download_imageset, + "_resolve_url", + Mock(return_value=download_imageset.UrlMeta( + suggested_filename=f"download{extension}", + suggested_extension=extension, + size_bytes=destination.stat().st_size, + )), + ) + run_command = Mock() + monkeypatch.setattr(download_imageset, "run_command", run_command) + monkeypatch.setattr(download_imageset, "process_imageset", Mock()) + + result = CliRunner().invoke( + cli=download_imageset.cli, + args=["filehost", url, "--name", "album", "--cwd", str(tmp_path), "--token", "test"], + ) + + assert result.exit_code == 0, result.output + if unfinished: + assert f"unfinished download; resuming {destination}" in caplog.text + run_command.assert_called_once_with(args=[ + "aria2c", "--continue=true", f"--dir={destination.parent}", + f"--out={destination.name}", url, + ]) + else: + assert "unfinished download; resuming" not in caplog.text + run_command.assert_not_called() + + +@pytest.mark.parametrize("process_directory", [False, True]) +def test_process_skips_unfinished_archive( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + process_directory: bool, +) -> None: + archive_path = tmp_path / "gallery.zip" + with zipfile.ZipFile(file=archive_path, mode="w") as archive: + archive.writestr(zinfo_or_arcname="photo.png", data=b"image data") + archive_path.with_name(f"{archive_path.name}.aria2").touch() + trash = Mock() + monkeypatch.setattr(download_imageset, "move_to_trash", trash) + + result = CliRunner().invoke( + cli=download_imageset.cli, + args=["process", str(tmp_path if process_directory else archive_path)], + ) + + assert result.exit_code == 0, result.output + assert "download is unfinished" in caplog.text + assert not (tmp_path / "photo.png").exists() + assert not (tmp_path / "gallery").exists() + assert archive_path.is_file() + trash.assert_not_called() + + if __name__ == "__main__": sys.exit(pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]]))