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

Author: pawank925

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)

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,29 @@ 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: dict[str, str] = dict(headers or {})
+    if user_agent:
+        base_headers = {
+            name: value
+            for name, value in base_headers.items()
+            if name.lower() != "user-agent"
+        }
+        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

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: