Use default User-Agent and accept headers in save_large_file (#360)

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)
 

src/zimscraperlib/download.py

@@ -165,112 +165,121 @@
         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,
     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.
 
     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 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
             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(

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"}