OCR PDF: Add missing languages body parameter

Assisted-by: Codex

Author: datalogics-cgreen

src/pdfrest/client.py

@@ -114,6 +114,7 @@
     GifColorModel,
     GraphicSmoothing,
     JpegColorModel,
+    OcrLanguage,
     PdfAType,
     PdfInfoQuery,
     PdfMergeInput,
@@ -2264,6 +2265,7 @@ def ocr_pdf(
         self,
         file: PdfRestFile | Sequence[PdfRestFile],
         *,
+        languages: OcrLanguage | Sequence[OcrLanguage] = "English",
         pages: PdfPageSelection | None = None,
         output: str | None = None,
         extra_query: Query | None = None,
@@ -2273,7 +2275,7 @@ def ocr_pdf(
     ) -> PdfRestFileBasedResponse:
         """Perform OCR on a PDF to make text searchable and extractable."""
 
-        payload: dict[str, Any] = {"files": file}
+        payload: dict[str, Any] = {"files": file, "languages": languages}
         if pages is not None:
             payload["pages"] = pages
         if output is not None:
@@ -3263,6 +3265,7 @@ async def ocr_pdf(
         self,
         file: PdfRestFile | Sequence[PdfRestFile],
         *,
+        languages: OcrLanguage | Sequence[OcrLanguage] = "English",
         pages: PdfPageSelection | None = None,
         output: str | None = None,
         extra_query: Query | None = None,
@@ -3272,7 +3275,7 @@ async def ocr_pdf(
     ) -> PdfRestFileBasedResponse:
         """Perform OCR on a PDF to make text searchable and extractable."""
 
-        payload: dict[str, Any] = {"files": file}
+        payload: dict[str, Any] = {"files": file, "languages": languages}
         if pages is not None:
             payload["pages"] = pages
         if output is not None:

src/pdfrest/models/_internal.py

@@ -22,6 +22,7 @@
 from pdfrest.types.public import PdfRedactionPreset
 
 from ..types import (
+    OcrLanguage,
     PdfAType,
     PdfInfoQuery,
     PdfXType,
@@ -329,6 +330,18 @@ class OcrPdfPayload(BaseModel):
         ),
         PlainSerializer(_serialize_as_first_file_id),
     ]
+    languages: Annotated[
+        list[OcrLanguage],
+        Field(
+            serialization_alias="languages",
+            validation_alias=AliasChoices("languages", "language"),
+            min_length=1,
+            default_factory=lambda: ["English"],
+        ),
+        BeforeValidator(_ensure_list),
+        BeforeValidator(_split_comma_list),
+        PlainSerializer(_serialize_as_comma_separated_string),
+    ]
     pages: Annotated[
         list[AscendingPageRange] | None,
         Field(serialization_alias="pages", min_length=1, default=None),

src/pdfrest/types/__init__.py

@@ -1,6 +1,7 @@
 """Public import surface for shared pdfrest types."""
 
 from .public import (
+    ALL_OCR_LANGUAGES,
     ALL_PDF_INFO_QUERIES,
     BmpColorModel,
     CompressionLevel,
@@ -9,6 +10,7 @@
     GifColorModel,
     GraphicSmoothing,
     JpegColorModel,
+    OcrLanguage,
     PdfAType,
     PdfInfoQuery,
     PdfMergeInput,
@@ -28,6 +30,7 @@
 )
 
 __all__ = [
+    "ALL_OCR_LANGUAGES",
     "ALL_PDF_INFO_QUERIES",
     "BmpColorModel",
     "CompressionLevel",
@@ -36,6 +39,7 @@
     "GifColorModel",
     "GraphicSmoothing",
     "JpegColorModel",
+    "OcrLanguage",
     "PdfAType",
     "PdfInfoQuery",
     "PdfMergeInput",

src/pdfrest/types/public.py

@@ -13,6 +13,7 @@
     PdfRestFile = Any
 
 __all__ = (
+    "ALL_OCR_LANGUAGES",
     "ALL_PDF_INFO_QUERIES",
     "BmpColorModel",
     "CompressionLevel",
@@ -21,6 +22,7 @@
     "GifColorModel",
     "GraphicSmoothing",
     "JpegColorModel",
+    "OcrLanguage",
     "PdfAType",
     "PdfInfoQuery",
     "PdfMergeInput",
@@ -140,3 +142,21 @@ class PdfMergeSource(TypedDict, total=False):
 SummaryOutputType = Literal["json", "file"]
 
 TranslateOutputFormat = Literal["plaintext", "markdown"]
+
+OcrLanguage = Literal[
+    "ChineseSimplified",
+    "ChineseTraditional",
+    "Dutch",
+    "English",
+    "French",
+    "German",
+    "Italian",
+    "Japanese",
+    "Korean",
+    "Portuguese",
+    "Spanish",
+]
+
+ALL_OCR_LANGUAGES: tuple[OcrLanguage, ...] = cast(
+    tuple[OcrLanguage, ...], get_args(OcrLanguage)
+)

tests/live/test_live_ocr_pdf.py

@@ -18,7 +18,7 @@ def test_live_ocr_pdf_success(
         base_url=pdfrest_live_base_url,
     ) as client:
         uploaded = client.files.create_from_paths([resource])[0]
-        response = client.ocr_pdf(uploaded)
+        response = client.ocr_pdf(uploaded, languages=["English", "German"])
 
     assert isinstance(response, PdfRestFileBasedResponse)
     assert response.output_files

tests/test_ocr_pdf.py

@@ -38,11 +38,36 @@ def test_ocr_payload_invalid_page_range() -> None:
         OcrPdfPayload.model_validate({"files": [file_repr], "pages": ["5-2"]})
 
 
+def test_ocr_payload_languages() -> None:
+    file_repr = make_pdf_file(PdfRestFileID.generate(1))
+    payload = OcrPdfPayload.model_validate(
+        {"files": [file_repr], "languages": ["English", "German"]}
+    )
+    assert payload.languages == ["English", "German"]
+    assert (
+        payload.model_dump(
+            mode="json", by_alias=True, exclude_none=True, exclude_unset=True
+        )["languages"]
+        == "English,German"
+    )
+
+
+def test_ocr_payload_invalid_language() -> None:
+    file_repr = make_pdf_file(PdfRestFileID.generate(1))
+    with pytest.raises(ValidationError, match="ChineseSimplified"):
+        OcrPdfPayload.model_validate({"files": [file_repr], "languages": ["Klingon"]})
+
+
 def test_ocr_pdf_success(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.delenv("PDFREST_API_KEY", raising=False)
     input_file = make_pdf_file(PdfRestFileID.generate(1))
     payload_dump = OcrPdfPayload.model_validate(
-        {"files": [input_file], "pages": ["1-3"], "output": "ocr"}
+        {
+            "files": [input_file],
+            "pages": ["1-3"],
+            "output": "ocr",
+            "languages": ["English"],
+        }
     ).model_dump(mode="json", by_alias=True, exclude_none=True, exclude_unset=True)
     output_id = str(PdfRestFileID.generate())
 
@@ -91,9 +116,9 @@ def test_ocr_pdf_request_customization(
 ) -> None:
     monkeypatch.delenv("PDFREST_API_KEY", raising=False)
     input_file = make_pdf_file(PdfRestFileID.generate(1))
-    payload_dump = OcrPdfPayload.model_validate({"files": [input_file]}).model_dump(
-        mode="json", by_alias=True, exclude_none=True, exclude_unset=True
-    )
+    payload_dump = OcrPdfPayload.model_validate(
+        {"files": [input_file], "languages": ["English"]}
+    ).model_dump(mode="json", by_alias=True, exclude_none=True, exclude_unset=True)
     output_id = str(PdfRestFileID.generate())
     captured_timeout: dict[str, float | dict[str, float] | None] = {}
 
@@ -152,9 +177,9 @@ async def test_async_ocr_pdf_success(
 ) -> None:
     monkeypatch.delenv("PDFREST_API_KEY", raising=False)
     input_file = make_pdf_file(PdfRestFileID.generate(2))
-    payload_dump = OcrPdfPayload.model_validate({"files": [input_file]}).model_dump(
-        mode="json", by_alias=True, exclude_none=True, exclude_unset=True
-    )
+    payload_dump = OcrPdfPayload.model_validate(
+        {"files": [input_file], "languages": ["English"]}
+    ).model_dump(mode="json", by_alias=True, exclude_none=True, exclude_unset=True)
     output_id = str(PdfRestFileID.generate())
 
     seen: dict[str, int] = {"post": 0, "get": 0}