From b986c95d8efb3c6dac7c626f7dacd798586902d7 Mon Sep 17 00:00:00 2001 From: adityahabib906 Date: Thu, 22 Oct 2026 19:56:01 +0530 Subject: [PATCH] Use default User-Agent and accept headers in save_large_file (#360) Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01T9Jc2YLGF81qNgcakGKSgg --- CHANGELOG.md | 1 + src/zimscraperlib/download.py | 15 +++++-- tests/download/test_download.py | 70 +++++++++++++++++++++++++++++++++ 3 files changed, 83 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8a7c209..187f003 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,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) - Add fuzzy rule to rewrite URLs of lesbases.anct.gouv.fr (#334) +- Send a Wikimedia-compliant default `User-Agent` in `save_large_file` when downloading from `upload.wikimedia.org`, and accept arbitrary request headers through a new `headers` argument (#360) - 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..cf7a0b5 100644 --- a/src/zimscraperlib/download.py +++ b/src/zimscraperlib/download.py @@ -171,6 +171,7 @@ def save_large_file( 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,26 @@ 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, ignored if a + User-Agent is already set in `headers` on_progress: optional callback called with the bytes downloaded so far and the total size, when known + headers: additional headers to send with the requests. When downloading from + `upload.wikimedia.org` and no User-Agent is provided, a Wikimedia-compliant + default one is sent """ - headers: dict[str, str] = {"User-Agent": user_agent} if user_agent else {} + base_headers: dict[str, str] = dict(headers or {}) + if user_agent and not has_user_agent_header(base_headers): + base_headers["User-Agent"] = user_agent + if is_wikimedia_upload_url(url) and not has_user_agent_header(base_headers): + base_headers["User-Agent"] = DEFAULT_USER_AGENT 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 diff --git a/tests/download/test_download.py b/tests/download/test_download.py index 0ff2f33..731165c 100644 --- a/tests/download/test_download.py +++ b/tests/download/test_download.py @@ -684,3 +684,73 @@ def test_get_options_override_outtmpl_other_vars( assert overriden[key] == value else: assert overriden[key] == custom_outtmpl + + +WIKIMEDIA_FILE_URL = "https://upload.wikimedia.org/a/b/file.zim" + + +def test_save_large_file_wikimedia_default_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) + + save_large_file(WIKIMEDIA_FILE_URL, tmp_path / "out.zim") + + _args, kwargs = get_mock.call_args + assert kwargs["headers"] == {"User-Agent": DEFAULT_USER_AGENT} + + +@pytest.mark.parametrize("ua_key", ["User-Agent", "user-agent"]) +def test_save_large_file_wikimedia_custom_user_agent_header_kept( + tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch, ua_key: str +) -> None: + get_mock = MagicMock(return_value=FakeResponse(200, [b"data"])) + monkeypatch.setattr(requests.Session, "get", get_mock) + + save_large_file( + WIKIMEDIA_FILE_URL, tmp_path / "out.zim", headers={ua_key: "mine/1.0"} + ) + + _args, kwargs = get_mock.call_args + assert kwargs["headers"] == {ua_key: "mine/1.0"} + + +def test_save_large_file_wikimedia_user_agent_argument_kept( + tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch +) -> None: + get_mock = MagicMock(return_value=FakeResponse(200, [b"data"])) + monkeypatch.setattr(requests.Session, "get", get_mock) + + save_large_file(WIKIMEDIA_FILE_URL, tmp_path / "out.zim", user_agent="mine/1.0") + + _args, kwargs = get_mock.call_args + assert kwargs["headers"] == {"User-Agent": "mine/1.0"} + + +def test_save_large_file_extra_headers_merged( + tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch +) -> None: + get_mock = MagicMock(return_value=FakeResponse(200, [b"data"])) + monkeypatch.setattr(requests.Session, "get", get_mock) + extra = {"X-Test": "1"} + + save_large_file(WIKIMEDIA_FILE_URL, tmp_path / "out.zim", headers=extra) + + _args, kwargs = get_mock.call_args + assert kwargs["headers"] == {"X-Test": "1", "User-Agent": DEFAULT_USER_AGENT} + assert extra == {"X-Test": "1"} # caller's dict is not modified + + +def test_save_large_file_other_host_no_default_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) + + save_large_file( + "https://example.com/file.zim", tmp_path / "out.zim", headers={"X-Test": "1"} + ) + + _args, kwargs = get_mock.call_args + assert kwargs["headers"] == {"X-Test": "1"}