stream output

Author: jdkent

elsevier_coordinate_extraction/download/api.py

@@ -2,9 +2,10 @@
 
 from __future__ import annotations
 
+import inspect
 import re
-from collections.abc import Mapping, Sequence
-from typing import Any
+from collections.abc import Awaitable, Mapping, Sequence
+from typing import Any, Protocol
 
 import httpx
 from lxml import etree
@@ -43,20 +44,39 @@
 _CDN_BASE = "https://ars.els-cdn.com/content/image"
 
 
+class ProgressCallback(
+    Protocol,
+):
+    """Callback invoked after each record is processed."""
+
+    def __call__(
+        self,
+        record: Mapping[str, str],
+        article: ArticleContent | None,
+        error: BaseException | None,
+    ) -> Awaitable[None] | None: ...
+
+
 async def download_articles(
     records: Sequence[Mapping[str, str]],
     *,
     client: ScienceDirectClient | None = None,
     cache: Any | None = None,
     cache_namespace: str = "articles",
     settings: Settings | None = None,
+    progress_callback: ProgressCallback | None = None,
 ) -> list[ArticleContent]:
     """Download ScienceDirect articles identified by DOI and/or PubMed ID records.
 
     Each record in ``records`` should contain at least one of the keys ``"doi"`` or ``"pmid"``.
     For every record, the downloader first attempts to retrieve the FULL text using the DOI
     (when present); if that fails, it retries with the PubMed ID. A successful download using
     either identifier stops further attempts for that record.
+
+    When ``progress_callback`` is provided it will be invoked after each record finishes processing.
+    The callback receives the original record, the downloaded ``ArticleContent`` when successful
+    (``None`` when no payload is returned), and the exception raised while processing
+    (``None`` on success). Callbacks may be synchronous or async functions.
     """
     if not records:
         return []
@@ -65,23 +85,39 @@ async def download_articles(
     owns_client = client is None
     sci_client = client or ScienceDirectClient(cfg)
 
+    async def _emit_progress(
+        record: Mapping[str, str],
+        article: ArticleContent | None,
+        error: BaseException | None,
+    ) -> None:
+        if progress_callback is None:
+            return
+        result = progress_callback(record, article, error)
+        if inspect.isawaitable(result):
+            await result
+
     async def _runner() -> list[ArticleContent]:
         results: list[ArticleContent] = []
         for record in records:
+            article: ArticleContent | None = None
             try:
                 article = await _download_record(
                     record=record,
                     client=sci_client,
                     cache=cache,
                     cache_namespace=cache_namespace,
                 )
-            except httpx.HTTPError:
+            except httpx.HTTPError as exc:
+                await _emit_progress(record, None, exc)
                 raise
-            except Exception:
+            except Exception as exc:
+                await _emit_progress(record, None, exc)
                 continue
             if article is None:
+                await _emit_progress(record, None, None)
                 continue
             results.append(article)
+            await _emit_progress(record, article, None)
         return results
 
     if owns_client:

tests/download/test_api.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from collections.abc import Sequence
 from pathlib import Path
 
@@ -221,3 +222,80 @@ async def request(
     assert article.metadata.get("rate_limit_limit") == 100
     assert article.metadata.get("rate_limit_remaining") == 99
     assert article.metadata.get("rate_limit_reset_epoch") == 1234567891.0
+
+
+@pytest.mark.asyncio()
+async def test_download_progress_callback_invoked_for_each_record(test_dois: Sequence[str]) -> None:
+    """Progress callback fires for every successfully downloaded record."""
+
+    cfg = _test_settings()
+    records = [{"doi": test_dois[0]}, {"doi": test_dois[1]}]
+
+    async def handler(request: httpx.Request) -> httpx.Response:
+        doi = request.url.path.rsplit("/", 1)[-1]
+        payload = f"""
+        <article xmlns="http://www.elsevier.com/xml/svapi/article/dtd" xmlns:ce="http://www.elsevier.com/xml/common/dtd">
+          <item-info>
+            <doi>{doi}</doi>
+            <pii>S105381192400679X</pii>
+          </item-info>
+          <ce:body><ce:para>{doi}</ce:para></ce:body>
+        </article>
+        """.encode("utf-8")
+        return httpx.Response(
+            200,
+            content=payload,
+            headers={"content-type": "application/xml"},
+            request=request,
+        )
+
+    transport = httpx.MockTransport(handler)
+    progress_calls: list[tuple[dict[str, str], ArticleContent | None, BaseException | None]] = []
+
+    def progress_cb(
+        record: dict[str, str],
+        article: ArticleContent | None,
+        error: BaseException | None,
+    ) -> None:
+        progress_calls.append((record, article, error))
+
+    async with ScienceDirectClient(cfg, transport=transport) as client:
+        articles = await download_articles(records, client=client, progress_callback=progress_cb)
+
+    assert len(articles) == len(records)
+    assert len(progress_calls) == len(records)
+    assert [call[0]["doi"] for call in progress_calls] == [record["doi"] for record in records]
+    assert all(call[1] is not None for call in progress_calls)
+    assert all(call[2] is None for call in progress_calls)
+
+
+@pytest.mark.asyncio()
+async def test_download_progress_callback_receives_errors(test_dois: Sequence[str]) -> None:
+    """Progress callback should receive exceptions before they propagate."""
+
+    cfg = _test_settings()
+    doi = test_dois[0]
+
+    async def handler(request: httpx.Request) -> httpx.Response:
+        raise httpx.TimeoutException("simulated timeout")
+
+    transport = httpx.MockTransport(handler)
+    progress_calls: list[tuple[dict[str, str], ArticleContent | None, BaseException | None]] = []
+
+    async def progress_cb(
+        record: dict[str, str],
+        article: ArticleContent | None,
+        error: BaseException | None,
+    ) -> None:
+        await asyncio.sleep(0)
+        progress_calls.append((record, article, error))
+
+    with pytest.raises(httpx.TimeoutException):
+        async with ScienceDirectClient(cfg, transport=transport) as client:
+            await download_articles([{"doi": doi}], client=client, progress_callback=progress_cb)
+
+    assert len(progress_calls) == 1
+    record, article, error = progress_calls[0]
+    assert record["doi"] == doi
+    assert article is None
+    assert isinstance(error, httpx.TimeoutException)