From 6b90e6603d6a79c95fe9a8e1fab0a7600e5f3d2b Mon Sep 17 00:00:00 2001 From: pawank925 <229111348+pawank925@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:39:29 +0530 Subject: [PATCH] Send default User-Agent and accept headers in save_large_file save_large_file only accepted a user_agent parameter, so callers could not set other headers, and unlike stream_file it did not send the Wikimedia-compliant default User-Agent for upload.wikimedia.org. Accept arbitrary headers, merge them with the existing user_agent, and fall back to the default User-Agent for Wikimedia upload URLs when none is set. Closes #360 --- CHANGELOG.md | 1 + src/zimscraperlib/download.py | 30 +++++++++-- tests/download/test_download.py | 90 +++++++++++++++++++++++++++++++++ 3 files changed, 117 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e3036b4..ebce448 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ 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) +- `save_large_file` accepts arbitrary HTTP headers and sends the default Wikimedia-compliant `User-Agent` for `upload.wikimedia.org`, like `stream_file` (#360) - 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) - New `get_pdf_title` and `get_epub_title` helpers to get a document's title from its own metadata without extracting its full text content, cheaper than `get_pdf_index_data`/`get_epub_index_data` when only the title is needed (openzim/warc2zim#462) diff --git a/src/zimscraperlib/download.py b/src/zimscraperlib/download.py index 9bf0e46..4fb55f0 100644 --- a/src/zimscraperlib/download.py +++ b/src/zimscraperlib/download.py @@ -165,12 +165,13 @@ def _log_download_progress(url: str, downloaded: int, total_size: int | None) -> logger.info(f"Downloaded {_format_size(downloaded)} of {url}") -def save_large_file( +def save_large_file( # noqa: PLR0917 url: str, fpath: pathlib.Path, retries: int = 5, user_agent: str | None = None, on_progress: Callable[[int, int | None], None] | None = None, + headers: dict[str, str] | None = None, ) -> None: """Download a binary file from its URL, streaming it with resume support. @@ -189,18 +190,20 @@ def save_large_file( 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 + user_agent: User-Agent header to send with the requests, takes precedence + over a User-Agent entry in `headers` on_progress: optional callback called with the bytes downloaded so far and the total size, when known + headers: additional HTTP headers to send with the requests """ - headers: dict[str, str] = {"User-Agent": user_agent} if user_agent else {} + base_headers = _build_request_headers(url, user_agent, headers) 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) + request_headers = dict(base_headers) if resume_from: request_headers["Range"] = f"bytes={resume_from}-" bytes_written = 0 @@ -307,6 +310,25 @@ def has_user_agent_header(headers: dict[str, str] | None) -> bool: return any(key.lower() == "user-agent" for key in (headers or {})) +def _build_request_headers( + url: str, + user_agent: str | None, + headers: dict[str, str] | None, +) -> dict[str, str]: + """Merge explicit headers with user_agent and the Wikimedia default User-Agent""" + request_headers: dict[str, str] = dict(headers or {}) + if user_agent: + request_headers = { + name: value + for name, value in request_headers.items() + if name.lower() != "user-agent" + } + request_headers["User-Agent"] = user_agent + if is_wikimedia_upload_url(url) and not has_user_agent_header(request_headers): + request_headers["User-Agent"] = DEFAULT_USER_AGENT + return request_headers + + def stream_file( # noqa: PLR0917 url: str, fpath: pathlib.Path | None = None, diff --git a/tests/download/test_download.py b/tests/download/test_download.py index 0ff2f33..96d3d60 100644 --- a/tests/download/test_download.py +++ b/tests/download/test_download.py @@ -376,6 +376,96 @@ def test_save_large_file_without_user_agent( assert "User-Agent" not in kwargs["headers"] +def test_save_large_file_default_user_agent_for_wikimedia( + 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://upload.wikimedia.org/wikipedia/commons/4/4a/Commons-logo.svg", + fpath, + ) + + _args, kwargs = get_mock.call_args + assert kwargs["headers"]["User-Agent"] == DEFAULT_USER_AGENT + + +def test_save_large_file_accepts_arbitrary_headers( + 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, headers={"Accept": "image/png"} + ) + + _args, kwargs = get_mock.call_args + assert kwargs["headers"]["Accept"] == "image/png" + assert "User-Agent" not in kwargs["headers"] + + +def test_save_large_file_user_agent_overrides_header( + 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, + user_agent="custom-agent", + headers={"user-agent": "header-agent", "Accept": "image/png"}, + ) + + _args, kwargs = get_mock.call_args + assert kwargs["headers"] == { + "Accept": "image/png", + "User-Agent": "custom-agent", + } + + +def test_save_large_file_keeps_header_user_agent_for_wikimedia( + 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://upload.wikimedia.org/wikipedia/commons/4/4a/Commons-logo.svg", + fpath, + headers={"User-Agent": "header-agent"}, + ) + + _args, kwargs = get_mock.call_args + assert kwargs["headers"]["User-Agent"] == "header-agent" + + +def test_save_large_file_default_user_agent_for_wikimedia_with_other_headers( + 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://upload.wikimedia.org/wikipedia/commons/4/4a/Commons-logo.svg", + fpath, + headers={"Accept": "image/png"}, + ) + + _args, kwargs = get_mock.call_args + assert kwargs["headers"] == { + "Accept": "image/png", + "User-Agent": DEFAULT_USER_AGENT, + } + + def test_save_large_file_resumes_after_connection_error( tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch ) -> None: