Fixing jpg handling in ParsedMedia.to_image_url (#1150)

Author: jamesbraza

src/paperqa/types.py

@@ -549,8 +549,11 @@ def __eq__(self, other) -> bool:
             == json.dumps(other.info, sort_keys=True)
         )
 
-    def to_image_url(self, image_type: str = "png") -> str:
+    def to_image_url(self) -> str:
         """Convert the image data to an RFC 2397 data URL format."""
+        image_type = cast(str, self.info.get("suffix", "png")).removeprefix(".")
+        if image_type == "jpg":  # SEE: https://stackoverflow.com/a/54488403
+            image_type = "jpeg"
         return encode_image_as_url(image_type, self.data)
 
     def save(self, path: str | os.PathLike) -> None:

tests/test_paperqa.py

@@ -1448,6 +1448,28 @@ async def test_chunk_metadata_reader(
         )
 
 
+def test_media_to_image_url(subtests: SubTests) -> None:
+    with subtests.test(msg="jpg"):
+        media = ParsedMedia(index=0, data=b"fake_jpg", info={"suffix": ".jpg"})
+        url = media.to_image_url()
+        assert "image/jpeg" in url
+
+    with subtests.test(msg="jpeg"):
+        media = ParsedMedia(index=0, data=b"fake_jpeg", info={"suffix": ".jpeg"})
+        url = media.to_image_url()
+        assert "image/jpeg" in url
+
+    with subtests.test(msg="png"):
+        media = ParsedMedia(index=0, data=b"fake_png", info={"suffix": ".png"})
+        url = media.to_image_url()
+        assert "image/png" in url
+
+    with subtests.test(msg="default"):
+        media = ParsedMedia(index=0, data=b"fake_png")
+        url = media.to_image_url()
+        assert "image/png" in url
+
+
 @pytest.mark.asyncio
 async def test_image_aggregation(stub_data_dir: Path) -> None:
     png_path = stub_data_dir / "sf_districts.png"