Added validator arg to parse_image, and used in testing to confirm concurrency works

Author: jamesbraza

src/paperqa/readers.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import os
+from collections.abc import Awaitable, Callable
 from math import ceil
 from pathlib import Path
 from typing import Literal, Protocol, cast, overload, runtime_checkable
@@ -32,9 +33,18 @@ def __call__(
     ) -> ParsedText: ...
 
 
-async def parse_image(path: str | os.PathLike, **_) -> ParsedText:
+async def parse_image(
+    path: str | os.PathLike, validator: Callable[[bytes], Awaitable] | None = None, **_
+) -> ParsedText:
     apath = anyio.Path(path)
     image_data = await anyio.Path(path).read_bytes()
+    if validator:
+        try:
+            await validator(image_data)
+        except Exception as exc:
+            raise ImpossibleParsingError(
+                f"Image validation failed for the image at path {path}."
+            ) from exc
     parsed_media = ParsedMedia(index=0, data=image_data, info={"suffix": apath.suffix})
     metadata = ParsedMetadata(
         parsing_libraries=[],

tests/test_paperqa.py

@@ -14,6 +14,7 @@
 from io import BytesIO
 from pathlib import Path
 from typing import cast
+from unittest.mock import MagicMock, call
 from uuid import UUID
 
 import httpx
@@ -1427,13 +1428,31 @@ async def test_read_doc_images_metadata(stub_data_dir: Path) -> None:
 async def test_read_doc_images_concurrency(stub_data_dir: Path) -> None:
     png_path = stub_data_dir / "sf_districts.png"
     doc = Doc(docname="stub", citation="stub", dockey="stub")
+    validation_mock = MagicMock()
+
+    async def validate(data: bytes) -> None:  # noqa: RUF029
+        validate_image(io.BytesIO(data))
+        validation_mock(data)
 
     # Check we can concurrently read in the same image many times
-    bulk_texts = await asyncio.gather(*(read_doc(png_path, doc) for _ in range(10)))
+    concurrent_call_count = 10
+    seen_media = set()
+    bulk_texts = await asyncio.gather(
+        *(
+            read_doc(png_path, doc, validator=validate)
+            for _ in range(concurrent_call_count)
+        )
+    )
     for (text,) in bulk_texts:
         assert text.doc == doc
         assert len(text.media) == 1
-        validate_image(io.BytesIO(text.media[0].data))
+        seen_media.add(text.media[0])
+    assert (
+        len(seen_media) == 1
+    ), "Expected the concurrent reads to all have the same parsed result"
+    validation_mock.assert_has_calls(
+        [call(next(iter(seen_media)).data)] * concurrent_call_count
+    )
 
 
 @pytest.mark.asyncio