From 1353fa9b0d6eab975322325a18fd823eca166aaa Mon Sep 17 00:00:00 2001
From: pawank925 <229111348+pawank925@users.noreply.github.com>
Date: Wed, 14 Oct 2026 11:10:44 +0530
Subject: [PATCH] Make save_large_file resumable with progress-based retries

---
 CHANGELOG.md                    |   2 +
 src/zimscraperlib/download.py   | 173 ++++++++++++++++++----
 tests/download/test_download.py | 252 +++++++++++++++++++++++++++++++-
 3 files changed, 393 insertions(+), 34 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 41ff6fb..4e9e06b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,9 +12,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Index ePub documents for full-text search, with a new `get_epub_index_data` helper and automatic indexing of ePub items, like PDFs already are (#333)
 - Send a Wikimedia-compliant default `User-Agent` in `stream_file` when downloading from `upload.wikimedia.org` (#214)
 - Add fuzzy rule to rewrite URLs of lesbases.anct.gouv.fr (#334)
+- Report progress from `save_large_file` through a new `on_progress` callback and periodic logging, and stream the download to a `.part` file that is only moved into place once complete (#348)
 
 ### Fixed
 
+- Rewrite `save_large_file` to stream and resume the download itself instead of running `wget` once, whose `-t` counter is a budget for the whole download: a flaky connection now keeps resuming and the retry budget only counts attempts that make no progress, so long downloads are no longer abandoned mid-way (#348)
 - Carry a script's top-level `const`, `let` and `class` declarations across the wombat block, so other scripts on the page can still see them (#329)
 - Keep a rewritten `<script>`'s original `src` in `__wb_orig_src`, so wombat can return it from `getAttribute("src")` and bundlers that identify chunks by it (Next.js/Turbopack, Vite) still work (openzim/warc2zim#473)
 - Keep URL-encoded `#` characters of the path or querystring when rewriting document URIs, instead of dropping everything after them (#341)
diff --git a/src/zimscraperlib/download.py b/src/zimscraperlib/download.py
index 64bccf9..9bf0e46 100644
--- a/src/zimscraperlib/download.py
+++ b/src/zimscraperlib/download.py
@@ -1,5 +1,7 @@
 import pathlib
-import subprocess
+import random
+import time
+from collections.abc import Callable
 from concurrent.futures import Future, ThreadPoolExecutor
 from typing import Any, ClassVar
 from urllib.parse import urlparse
@@ -126,33 +128,149 @@ class BestMp4(YoutubeConfig):
     }
 
 
+_RETRYABLE_HTTP_STATUS_CODES = (413, 429, 500, 502, 503, 504)
+_HTTP_PARTIAL_CONTENT = 206
+_DOWNLOAD_BLOCK_SIZE = 1024 * 1024
+_DOWNLOAD_CONNECT_TIMEOUT = 30
+_DOWNLOAD_READ_TIMEOUT = 60
+_DOWNLOAD_BACKOFF_FACTOR = 2
+_DOWNLOAD_BACKOFF_MAX = 60
+_DOWNLOAD_PROGRESS_INTERVAL = 30
+_SIZE_UNITS = ("B", "KiB", "MiB", "GiB", "TiB")
+_SIZE_STEP = 1024
+
+
+class DownloadError(Exception):
+    """Raised when a file could not be completely downloaded."""
+
+
+def _format_size(num_bytes: int) -> str:
+    """A byte count in a human-friendly binary unit, e.g. "1.5 GiB"."""
+    size = float(num_bytes)
+    for unit in _SIZE_UNITS[:-1]:
+        if size < _SIZE_STEP:
+            return f"{int(size)} B" if unit == "B" else f"{size:.1f} {unit}"
+        size /= _SIZE_STEP
+    return f"{size:.1f} {_SIZE_UNITS[-1]}"
+
+
+def _log_download_progress(url: str, downloaded: int, total_size: int | None) -> None:
+    """Log how much of the file has been downloaded so far."""
+    if total_size:
+        logger.info(
+            f"Downloaded {_format_size(downloaded)} ({downloaded / total_size:.0%}) "
+            f"of {url}"
+        )
+    else:
+        logger.info(f"Downloaded {_format_size(downloaded)} of {url}")
+
+
 def save_large_file(
-    url: str, fpath: pathlib.Path, retries: int = 5, user_agent: str | None = None
+    url: str,
+    fpath: pathlib.Path,
+    retries: int = 5,
+    user_agent: str | None = None,
+    on_progress: Callable[[int, int | None], None] | None = None,
 ) -> None:
-    """download a binary file from its URL, using wget
+    """Download a binary file from its URL, streaming it with resume support.
+
+    Unlike a single wget run, whose -t counter is a budget for the whole download,
+    the retry counter here only grows on attempts that make no progress at all:
+    any attempt that writes at least one byte resets it, so a download that keeps
+    resuming is not abandoned early.
+
+    The file is streamed to a sibling `<fpath>.part` file and only moved to `fpath`
+    once complete, so a watcher of the destination folder never sees a half-written
+    file. Progress is logged at most once every `_DOWNLOAD_PROGRESS_INTERVAL`
+    seconds and reported at the same pace through `on_progress` when given, with the
+    number of bytes downloaded so far and the total size when known.
 
     Arguments -
-        url:
+        url: URL of the file to download
+        fpath: path where the downloaded file is written
+        retries: number of consecutive attempts without progress tolerated
+        user_agent: User-Agent header to send with the requests
+        on_progress: optional callback called with the bytes downloaded so far and
+        the total size, when known
     """
-    command = [
-        "/usr/bin/env",
-        "wget",
-        "-t",
-        f"{retries}",
-        "--retry-connrefused",
-        "--random-wait",
-        "--progress=dot:giga",
-        "-O",
-        str(fpath),
-        "-c",
-        url,
-    ]
-    if user_agent:
-        command += ["-U", user_agent]
-    subprocess.run(
-        command,
-        check=True,
-    )
+    headers: dict[str, str] = {"User-Agent": user_agent} if user_agent else {}
+    attempts_without_progress = 0
+    tmp_fpath = fpath.with_name(fpath.name + ".part")
+
+    with requests.Session() as session:
+        while True:
+            resume_from = tmp_fpath.stat().st_size if tmp_fpath.exists() else 0
+            request_headers = dict(headers)
+            if resume_from:
+                request_headers["Range"] = f"bytes={resume_from}-"
+            bytes_written = 0
+            try:
+                with session.get(
+                    url,
+                    stream=True,
+                    headers=request_headers,
+                    timeout=(_DOWNLOAD_CONNECT_TIMEOUT, _DOWNLOAD_READ_TIMEOUT),
+                ) as response:
+                    if resume_from and response.status_code != _HTTP_PARTIAL_CONTENT:
+                        logger.warning(
+                            f"{url} does not support resuming, restarting download"
+                        )
+                        tmp_fpath.unlink(missing_ok=True)
+                        resume_from = 0
+                    response.raise_for_status()
+                    content_length = response.headers.get("Content-Length")
+                    total_size = (
+                        resume_from + int(content_length)
+                        if content_length is not None
+                        else None
+                    )
+                    last_logged_at = time.monotonic()
+                    with open(tmp_fpath, "ab" if resume_from else "wb") as fh:
+                        for chunk in response.iter_content(_DOWNLOAD_BLOCK_SIZE):
+                            if not chunk:
+                                continue
+                            fh.write(chunk)
+                            bytes_written += len(chunk)
+                            now = time.monotonic()
+                            if now - last_logged_at >= _DOWNLOAD_PROGRESS_INTERVAL:
+                                downloaded = resume_from + bytes_written
+                                _log_download_progress(url, downloaded, total_size)
+                                if on_progress is not None:
+                                    on_progress(downloaded, total_size)
+                                last_logged_at = now
+            except requests.HTTPError as exc:
+                status = exc.response.status_code if exc.response is not None else None
+                if status not in _RETRYABLE_HTTP_STATUS_CODES:
+                    raise DownloadError(
+                        f"Failed to download {url}: HTTP status {status}"
+                    ) from exc
+            except (
+                requests.ConnectionError,
+                requests.Timeout,
+                requests.exceptions.ChunkedEncodingError,
+            ) as exc:
+                logger.warning(f"Download of {url} interrupted ({exc})")
+            else:
+                tmp_fpath.replace(fpath)
+                return
+
+            if bytes_written:
+                attempts_without_progress = 0
+            else:
+                attempts_without_progress += 1
+            if attempts_without_progress > retries:
+                raise DownloadError(
+                    f"Failed to download {url}: {retries} attempts without progress"
+                )
+            wait = min(
+                _DOWNLOAD_BACKOFF_MAX,
+                _DOWNLOAD_BACKOFF_FACTOR * 2 ** (attempts_without_progress - 1),
+            )
+            wait += random.uniform(0, wait * 0.1)  # noqa: S311
+            logger.warning(
+                f"Retrying download of {url} in {wait:.0f}s from byte {resume_from}"
+            )
+            time.sleep(wait)
 
 
 def get_retry_adapter(
@@ -166,14 +284,7 @@ def get_retry_adapter(
         status=2,  # failure HTTP status (only those bellow)
         redirect=False,  # don't fail on redirections
         backoff_factor=30,  # sleep factor between retries
-        status_forcelist=[
-            413,
-            429,
-            500,
-            502,
-            503,
-            504,
-        ],  # force retry on the following codes
+        status_forcelist=list(_RETRYABLE_HTTP_STATUS_CODES),
     )
 
     return requests.adapters.HTTPAdapter(max_retries=retries)
diff --git a/tests/download/test_download.py b/tests/download/test_download.py
index fe4b001..0ff2f33 100644
--- a/tests/download/test_download.py
+++ b/tests/download/test_download.py
@@ -2,19 +2,22 @@
 import io
 import pathlib
 import re
+from collections.abc import Iterator
 from typing import ClassVar
-from unittest.mock import Mock
+from unittest.mock import MagicMock, Mock
 
 import pytest
 import requests
 import requests.structures
-from yt_dlp.utils import DownloadError
+from yt_dlp.utils import DownloadError as YtDlpDownloadError
 
 from zimscraperlib.constants import DEFAULT_USER_AGENT, DEFAULT_WEB_REQUESTS_TIMEOUT
 from zimscraperlib.download import (
     BestMp4,
     BestWebm,
+    DownloadError,
     YoutubeDownloader,
+    _format_size,  # pyright: ignore[reportPrivateUsage]
     has_user_agent_header,
     is_wikimedia_upload_url,
     save_large_file,
@@ -311,6 +314,249 @@ def test_large_download_https_custom_ua(tmp_path: pathlib.Path, valid_https_url:
     assert_downloaded_file(valid_https_url, dest_file)
 
 
+class FakeResponse:
+    """Minimal stand-in for requests.Response, usable as a context manager."""
+
+    def __init__(
+        self,
+        status_code: int,
+        chunks: list[bytes] | None = None,
+        raise_while_streaming: Exception | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> None:
+        self.status_code = status_code
+        self._chunks = chunks or []
+        self._raise_while_streaming = raise_while_streaming
+        self.headers = headers or {}
+
+    def __enter__(self) -> FakeResponse:
+        return self
+
+    def __exit__(self, *_args: object) -> None:
+        return None
+
+    def raise_for_status(self) -> None:
+        if self.status_code >= 400:
+            response = requests.Response()
+            response.status_code = self.status_code
+            raise requests.HTTPError(response=response)
+
+    def iter_content(self, _chunk_size: int) -> Iterator[bytes]:
+        yield from self._chunks
+        if self._raise_while_streaming:
+            raise self._raise_while_streaming
+
+
+def test_save_large_file_single_attempt(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A clean download writes all chunks in one attempt, without a Range header."""
+    get_mock = MagicMock(return_value=FakeResponse(200, [b"hello ", b"", b"world"]))
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+
+    fpath = tmp_path / "out.zim"
+    save_large_file("https://example.com/file.zim", fpath, user_agent="some-agent")
+
+    assert fpath.read_bytes() == b"hello world"
+    _args, kwargs = get_mock.call_args
+    assert kwargs["headers"] == {"User-Agent": "some-agent"}
+    assert "Range" not in kwargs["headers"]
+
+
+def test_save_large_file_without_user_agent(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    get_mock = MagicMock(return_value=FakeResponse(200, [b"data"]))
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+
+    fpath = tmp_path / "out.zim"
+    save_large_file("https://example.com/file.zim", fpath)
+
+    _args, kwargs = get_mock.call_args
+    assert "User-Agent" not in kwargs["headers"]
+
+
+def test_save_large_file_resumes_after_connection_error(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A dropped connection resumes with a Range header for the missing bytes."""
+    responses = [
+        FakeResponse(
+            200,
+            [b"hello "],
+            raise_while_streaming=requests.ConnectionError("dropped"),
+        ),
+        FakeResponse(206, [b"world"]),
+    ]
+    get_mock = MagicMock(side_effect=responses)
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+    monkeypatch.setattr("zimscraperlib.download.time.sleep", MagicMock())
+
+    fpath = tmp_path / "out.zim"
+    save_large_file("https://example.com/file.zim", fpath)
+
+    assert fpath.read_bytes() == b"hello world"
+    assert get_mock.call_count == 2
+    assert get_mock.call_args_list[1].kwargs["headers"]["Range"] == "bytes=6-"
+
+
+def test_save_large_file_retries_on_retryable_http_status(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    get_mock = MagicMock(side_effect=[FakeResponse(503), FakeResponse(200, [b"ok"])])
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+    monkeypatch.setattr("zimscraperlib.download.time.sleep", MagicMock())
+
+    fpath = tmp_path / "out.zim"
+    save_large_file("https://example.com/file.zim", fpath)
+
+    assert fpath.read_bytes() == b"ok"
+    assert get_mock.call_count == 2
+
+
+def test_save_large_file_restarts_when_server_ignores_range(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """If the server replies 200 (not 206) to a Range request, start over."""
+    fpath = tmp_path / "out.zim"
+    fpath.with_name("out.zim.part").write_bytes(b"stale-partial-data")
+
+    get_mock = MagicMock(return_value=FakeResponse(200, [b"fresh"]))
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+
+    save_large_file("https://example.com/file.zim", fpath)
+
+    assert fpath.read_bytes() == b"fresh"
+    assert not fpath.with_name("out.zim.part").exists()
+    assert get_mock.call_args.kwargs["headers"]["Range"] == "bytes=18-"
+
+
+def test_save_large_file_does_not_retry_non_retryable_http_status(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    get_mock = MagicMock(return_value=FakeResponse(404))
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+
+    fpath = tmp_path / "out.zim"
+    with pytest.raises(DownloadError):
+        save_large_file("https://example.com/file.zim", fpath)
+
+    assert get_mock.call_count == 1
+
+
+def test_save_large_file_gives_up_after_consecutive_failures_without_progress(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    get_mock = MagicMock(
+        side_effect=requests.ConnectionError("always fails before any byte")
+    )
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+    monkeypatch.setattr("zimscraperlib.download.time.sleep", MagicMock())
+
+    fpath = tmp_path / "out.zim"
+    with pytest.raises(DownloadError):
+        save_large_file("https://example.com/file.zim", fpath, retries=2)
+
+    assert get_mock.call_count == 3
+
+
+def test_save_large_file_keeps_retrying_as_long_as_progress_is_made(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Every attempt but the last writes a byte before failing; progress keeps
+    resetting the failure counter, so more than `retries` attempts in a row are
+    allowed as long as each one makes progress."""
+    failing_responses = [
+        FakeResponse(
+            200 if index == 0 else 206,
+            [b"x"],
+            raise_while_streaming=requests.ConnectionError("dropped"),
+        )
+        for index in range(4)
+    ]
+    get_mock = MagicMock(side_effect=[*failing_responses, FakeResponse(206, [b"x"])])
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+    monkeypatch.setattr("zimscraperlib.download.time.sleep", MagicMock())
+
+    fpath = tmp_path / "out.zim"
+    save_large_file("https://example.com/file.zim", fpath, retries=2)
+
+    assert fpath.read_bytes() == b"xxxxx"
+    assert get_mock.call_count == 5
+
+
+def test_save_large_file_leaves_no_part_file_on_success(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(
+        requests.Session, "get", MagicMock(return_value=FakeResponse(200, [b"data"]))
+    )
+
+    fpath = tmp_path / "out.zim"
+    save_large_file("https://example.com/file.zim", fpath)
+
+    assert fpath.read_bytes() == b"data"
+    assert not fpath.with_name("out.zim.part").exists()
+
+
+def test_save_large_file_reports_progress_with_total_size(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Progress is reported once the interval elapses, with the total size known."""
+    monkeypatch.setattr("zimscraperlib.download._DOWNLOAD_PROGRESS_INTERVAL", 0)
+    monkeypatch.setattr(
+        requests.Session,
+        "get",
+        MagicMock(
+            return_value=FakeResponse(
+                200,
+                [b"hello ", b"world"],
+                headers={"Content-Length": "11"},
+            )
+        ),
+    )
+    progress: list[tuple[int, int | None]] = []
+
+    fpath = tmp_path / "out.zim"
+    save_large_file(
+        "https://example.com/file.zim",
+        fpath,
+        on_progress=lambda downloaded, total: progress.append((downloaded, total)),
+    )
+
+    assert fpath.read_bytes() == b"hello world"
+    assert progress == [(6, 11), (11, 11)]
+
+
+def test_save_large_file_reports_progress_without_total_size(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Without a Content-Length header, progress is reported with an unknown total."""
+    monkeypatch.setattr("zimscraperlib.download._DOWNLOAD_PROGRESS_INTERVAL", 0)
+    monkeypatch.setattr(
+        requests.Session, "get", MagicMock(return_value=FakeResponse(200, [b"data"]))
+    )
+
+    fpath = tmp_path / "out.zim"
+    save_large_file("https://example.com/file.zim", fpath)
+
+    assert fpath.read_bytes() == b"data"
+
+
+@pytest.mark.parametrize(
+    "num_bytes,expected",
+    [
+        (0, "0 B"),
+        (500, "500 B"),
+        (1536, "1.5 KiB"),
+        (1024 * 1024, "1.0 MiB"),
+        (1024**4, "1.0 TiB"),
+    ],
+)
+def test_format_size(num_bytes: int, expected: str) -> None:
+    assert _format_size(num_bytes) == expected
+
+
 @pytest.mark.slow
 @pytest.mark.parametrize(
     "url,video_id",
@@ -353,7 +599,7 @@ def test_youtube_download_nowait(tmp_path: pathlib.Path):
 @pytest.mark.slow
 def test_youtube_download_error():
     yt_downloader = YoutubeDownloader(threads=1)
-    with pytest.raises(DownloadError, match=re.escape("is not a valid URL")):
+    with pytest.raises(YtDlpDownloadError, match=re.escape("is not a valid URL")):
         yt_downloader.download("11", BestMp4.get_options())
     yt_downloader.shutdown()