✨ add support for local loading

Author: sebastianMindee

mindee/client.py

@@ -4,6 +4,7 @@
 
 from mindee.error.mindee_error import MindeeClientError, MindeeError
 from mindee.error.mindee_http_error import handle_error
+from mindee.input import LocalResponse
 from mindee.input.page_options import PageOptions
 from mindee.input.sources import (
     Base64Input,
@@ -178,6 +179,23 @@ def enqueue(
             cropper,
         )
 
+    def load_prediction(
+        self, product_class: Type[Inference], local_response: LocalResponse
+    ) -> Union[AsyncPredictResponse, PredictResponse]:
+        """
+        Load a prediction.
+
+        :param product_class: Class of the product to use.
+        :param local_response: Local response to load.
+        :return:
+        """
+        try:
+            if local_response.json.get("job"):
+                return AsyncPredictResponse(product_class, local_response.json)
+            return PredictResponse(product_class, local_response.json)
+        except KeyError as exc:
+            raise MindeeError("No prediction found in local response.") from exc
+
     def parse_queued(
         self,
         product_class: Type[Inference],

mindee/input/__init__.py

@@ -1,3 +1,4 @@
+from mindee.input.local_response import LocalResponse
 from mindee.input.page_options import PageOptions
 from mindee.input.sources import (
     Base64Input,

mindee/input/local_response.py

@@ -0,0 +1,29 @@
+import json
+from pathlib import Path
+from typing import Union, BinaryIO, Dict, Any
+import io
+
+from mindee.error import MindeeError
+
+
+class LocalResponse:
+    json: Dict[str, Any]
+
+    def __init__(self, input_file: Union[BinaryIO, str, Path, bytes]):
+        input_binary: BinaryIO
+        if isinstance(input_file, BinaryIO):
+            input_binary = input_file
+            input_binary.seek(0)
+        elif isinstance(input_file, str) or isinstance(input_file, Path):
+            with open(input_file, 'rb') as f:
+                input_binary = io.BytesIO(f.read())
+        elif isinstance(input_file, bytes):
+            input_binary = io.BytesIO(input_file)
+        else:
+            raise TypeError('Incompatible type for input.')
+        try:
+            input_binary.seek(0)
+            self.json = json.load(input_binary)
+            input_binary.close()
+        except json.decoder.JSONDecodeError as exc:
+            raise MindeeError('File is not a valid dictionary.') from exc

tests/test_client.py

@@ -2,13 +2,16 @@
 
 import pytest
 
-from mindee import Client, PageOptions, product
-from mindee.error.mindee_error import MindeeClientError
+from mindee import AsyncPredictResponse, Client, PageOptions, PredictResponse, product
+from mindee.error.mindee_error import MindeeClientError, MindeeError
 from mindee.error.mindee_http_error import MindeeHTTPError
+from mindee.input import LocalResponse
 from mindee.input.sources import LocalInputSource
+from mindee.product import InternationalIdV2, InvoiceV4
 from mindee.product.invoice_splitter.invoice_splitter_v1 import InvoiceSplitterV1
 from mindee.product.receipt.receipt_v4 import ReceiptV4
-from tests.test_inputs import FILE_TYPES_DIR
+from tests.mindee_http.test_error import ERROR_DATA_DIR
+from tests.test_inputs import FILE_TYPES_DIR, PRODUCT_DATA_DIR
 from tests.utils import clear_envvars, dummy_envvars
 
 
@@ -113,3 +116,40 @@ def test_async_wrong_polling_delay(dummy_client: Client):
     input_doc = dummy_client.source_from_path(FILE_TYPES_DIR / "pdf" / "blank.pdf")
     with pytest.raises(MindeeClientError):
         dummy_client.enqueue_and_parse(InvoiceSplitterV1, input_doc, delay_sec=0)
+
+
+def test_local_response_from_sync_json(dummy_client: Client):
+    input_file = LocalResponse(
+        PRODUCT_DATA_DIR / "invoices" / "response_v4" / "complete.json"
+    )
+    with open(PRODUCT_DATA_DIR / "invoices" / "response_v4" / "summary_full.rst") as f:
+        reference_doc = f.read()
+    result = dummy_client.load_prediction(InvoiceV4, input_file)
+    assert isinstance(result, PredictResponse)
+    assert str(result.document) == reference_doc
+
+
+def test_local_response_from_async_json(dummy_client: Client):
+    input_file = LocalResponse(
+        PRODUCT_DATA_DIR / "international_id" / "response_v2" / "complete.json"
+    )
+    with open(
+        PRODUCT_DATA_DIR / "international_id" / "response_v2" / "summary_full.rst"
+    ) as f:
+        reference_doc = f.read()
+    result = dummy_client.load_prediction(InternationalIdV2, input_file)
+    assert isinstance(result, AsyncPredictResponse)
+    assert str(result.document) == reference_doc
+
+
+def test_local_response_from_invalid_file(dummy_client: Client):
+    with pytest.raises(MindeeError):
+        LocalResponse(
+            PRODUCT_DATA_DIR / "invoices" / "response_v4" / "summary_full.rst"
+        )
+
+
+def test_local_response_from_invalid_dict(dummy_client: Client):
+    input_file = LocalResponse(ERROR_DATA_DIR / "error_400_no_details.json")
+    with pytest.raises(MindeeError):
+        dummy_client.load_prediction(InvoiceV4, input_file)