✨ Use generic parse method, allowing for document versions

Author: ianardee

.pylintrc

@@ -13,15 +13,34 @@ pip install mindee
 
 Finally, Python away!
 
-### Off-the-Shelf Document
+### Off-the-Shelf Documents
+World-wide documents:
 ```python
-from mindee import Client
+from mindee import Client, documents
 
 # Init a new client
 mindee_client = Client(api_key="my-api-key")
 
-# Load a file from disk and parse it as an invoice
-api_response = mindee_client.doc_from_path("/path/to/the/invoice.pdf").parse("invoice")
+# Load a file from disk
+input_doc = mindee_client.doc_from_path("/path/to/the/invoice.pdf")
+# Parse the document as an invoice by passing the documents.Invoice type
+api_response = input_doc.parse(documents.TypeInvoiceV3)
+
+# Print a brief summary of the parsed data
+print(api_response.document)
+```
+
+Region-specific documents:
+```python
+from mindee import Client, documents
+
+# Init a new client
+mindee_client = Client(api_key="my-api-key")
+
+# Load a file from disk
+input_doc = mindee_client.doc_from_path("/path/to/the/check.jpg")
+# Parse the document as an invoice by passing the documents.us.TypeBankCheckV1 type
+api_response = input_doc.parse(documents.us.TypeBankCheckV1)
 
 # Print a brief summary of the parsed data
 print(api_response.document)
@@ -30,7 +49,7 @@ print(api_response.document)
 ### Custom Document (API Builder)
 
 ```python
-from mindee import Client
+from mindee import Client, documents
 
 # Init a new client and add your custom endpoint (document)
 mindee_client = Client(api_key="my-api-key").add_endpoint(
@@ -39,7 +58,9 @@ mindee_client = Client(api_key="my-api-key").add_endpoint(
 )
 
 # Load a file from disk and parse it
-api_response = mindee_client.doc_from_path("/path/to/the/card.jpg").parse("wnine")
+api_response = mindee_client.doc_from_path(
+    "/path/to/the/w9.jpg"
+).parse(documents.TypeCustomV1, "wnine")
 
 # Print a brief summary of the parsed data
 print(api_response.document)

README.md

@@ -1,32 +1,43 @@
 import argparse
 import json
 from argparse import Namespace
-from typing import Dict, Union
+from dataclasses import dataclass
+from typing import Dict, Generic, TypeVar
 
-from mindee import Client, PageOptions
+from mindee import Client, PageOptions, documents
 from mindee.client import DocumentClient
-from mindee.documents.base import serialize_for_json
-
-DOCUMENTS: Dict[str, Dict[str, Union[list, str]]] = {
-    "invoice": {
-        "help": "Invoice",
-        "doc_type": "invoice",
-    },
-    "receipt": {
-        "help": "Expense Receipt",
-        "doc_type": "receipt",
-    },
-    "passport": {
-        "help": "Passport",
-        "doc_type": "passport",
-    },
-    "financial": {
-        "help": "Financial Document (receipt or invoice)",
-        "doc_type": "financial_doc",
-    },
-    "custom": {
-        "help": "Custom document type from API builder",
-    },
+from mindee.documents.base import Document, serialize_for_json
+
+TypeDoc = TypeVar("TypeDoc", bound=Document)
+
+
+@dataclass
+class CommandConfig(Generic[TypeDoc]):
+    help: str
+    doc_class: TypeDoc
+
+
+DOCUMENTS: Dict[str, CommandConfig] = {
+    "invoice": CommandConfig(
+        help="Invoice",
+        doc_class=documents.TypeInvoiceV3,
+    ),
+    "receipt": CommandConfig(
+        help="Expense Receipt",
+        doc_class=documents.TypeReceiptV3,
+    ),
+    "passport": CommandConfig(
+        help="Passport",
+        doc_class=documents.TypePassportV1,
+    ),
+    "financial": CommandConfig(
+        help="Financial Document (receipt or invoice)",
+        doc_class=documents.TypeFinancialDocument,
+    ),
+    "custom": CommandConfig(
+        help="Custom document type from API builder",
+        doc_class=documents.TypeCustomV1,
+    ),
 }
 
 
@@ -51,10 +62,8 @@ def call_endpoint(args: Namespace):
             endpoint_name=args.doc_type,
             account_name=args.username,
         )
-        doc_type = args.doc_type
-    else:
-        info = DOCUMENTS[args.product_name]
-        doc_type = info["doc_type"]
+    info = DOCUMENTS[args.product_name]
+    doc_class = info.doc_class
 
     input_doc = _get_input_doc(client, args)
 
@@ -64,20 +73,23 @@ def call_endpoint(args: Namespace):
         page_options = None
     if args.product_name == "custom":
         parsed_data = input_doc.parse(
-            doc_type, username=args.username, page_options=page_options
+            doc_class,
+            endpoint_name=args.doc_type,
+            account_name=args.username,
+            page_options=page_options,
         )
     else:
         parsed_data = input_doc.parse(
-            doc_type, include_words=args.include_words, page_options=page_options
+            doc_class, include_words=args.include_words, page_options=page_options
         )
 
     if args.output_type == "raw":
         print(json.dumps(parsed_data.http_response, indent=2))
     elif args.output_type == "parsed":
-        doc = getattr(parsed_data, doc_type)
+        doc = parsed_data.document
         print(json.dumps(doc, indent=2, default=serialize_for_json))
     else:
-        print(getattr(parsed_data, doc_type))
+        print(parsed_data.document)
 
 
 def _parse_args() -> Namespace:
@@ -95,7 +107,7 @@ def _parse_args() -> Namespace:
         required=True,
     )
     for name, info in DOCUMENTS.items():
-        subp = subparsers.add_parser(name, help=info["help"])
+        subp = subparsers.add_parser(name, help=info.help)
         if name == "custom":
             subp.add_argument(
                 "-u",

mindee/__main__.py

@@ -1,13 +1,14 @@
 import json
-from typing import BinaryIO, Dict, Optional
+from typing import BinaryIO, Dict, Optional, Type
 
-from mindee.document_config import DocumentConfig, DocumentConfigDict
-from mindee.documents.bank_check import BankCheck
-from mindee.documents.custom_document import CustomDocument
+from mindee.documents.base import Document, TypeDocument
+from mindee.documents.config import DocumentConfig, DocumentConfigDict
+from mindee.documents.custom.custom_v1 import CustomV1
 from mindee.documents.financial_document import FinancialDocument
-from mindee.documents.invoice import Invoice
-from mindee.documents.passport import Passport
-from mindee.documents.receipt import Receipt
+from mindee.documents.invoice.invoice_v3 import InvoiceV3
+from mindee.documents.passport.passport_v1 import PassportV1
+from mindee.documents.receipt.receipt_v3 import ReceiptV3
+from mindee.documents.us.bank_check.bank_check_v1 import BankCheckV1
 from mindee.endpoints import OTS_OWNER, CustomEndpoint, HTTPException, StandardEndpoint
 from mindee.input.page_options import PageOptions
 from mindee.input.sources import (
@@ -21,6 +22,11 @@
 from mindee.response import PredictResponse
 
 
+def get_type_var_name(type_var) -> str:
+    """Get the name of the bound class."""
+    return type_var.__bound__.__name__
+
+
 class DocumentClient:
     input_doc: InputSource
     doc_configs: DocumentConfigDict
@@ -38,42 +44,49 @@ def __init__(
 
     def parse(
         self,
-        document_type: str,
-        username: Optional[str] = None,
+        document_class: TypeDocument,
+        endpoint_name: Optional[str] = None,
+        account_name: Optional[str] = None,
         include_words: bool = False,
         close_file: bool = True,
         page_options: Optional[PageOptions] = None,
-    ) -> PredictResponse:
+    ) -> PredictResponse[TypeDocument]:
         """
         Call prediction API on the document and parse the results.
 
-        :param document_type: Document type to parse
-        :param username: API username, the endpoint owner
+        :type document_class: DocT
+        :param endpoint_name: Document type to parse
+        :param account_name: API username, the endpoint owner
         :param include_words: Include all the words of the document in the response
         :param close_file: Whether to `close()` the file after parsing it.
             Set to `False` if you need to access the file after this operation.
         :param page_options: PageOptions object for cutting multipage documents.
         """
-        logger.debug("Parsing document as '%s'", document_type)
+        if get_type_var_name(document_class) != CustomV1.__name__:
+            endpoint_name = get_type_var_name(document_class)
+        elif endpoint_name is None:
+            raise RuntimeError("document_type is required for CustomDocument")
+
+        logger.debug("Parsing document as '%s'", endpoint_name)
 
         found = []
         for k in self.doc_configs.keys():
-            if k[1] == document_type:
+            if k[1] == endpoint_name:
                 found.append(k)
 
         if len(found) == 0:
-            raise RuntimeError(f"Document type not configured: {document_type}")
+            raise RuntimeError(f"Document type not configured: {endpoint_name}")
 
-        if username:
-            config_key = (username, document_type)
+        if account_name:
+            config_key = (account_name, endpoint_name)
         elif len(found) == 1:
             config_key = found[0]
         else:
             usernames = [k[0] for k in found]
             raise RuntimeError(
                 (
                     "Duplicate configuration detected.\n"
-                    f"You specified a document_type '{document_type}' in your custom config.\n"
+                    f"You specified a document_type '{endpoint_name}' in your custom config.\n"
                     "To avoid confusion, please add the 'account_name' attribute to "
                     f"the parse method, one of {usernames}."
                 )
@@ -87,11 +100,18 @@ def parse(
                 page_options.on_min_pages,
                 page_options.page_indexes,
             )
-        return self._make_request(doc_config, include_words, close_file)
+        return self._make_request(document_class, doc_config, include_words, close_file)
 
     def _make_request(
-        self, doc_config: DocumentConfig, include_words: bool, close_file: bool
-    ) -> PredictResponse:
+        self,
+        document_class: TypeDocument,
+        doc_config: DocumentConfig,
+        include_words: bool,
+        close_file: bool,
+    ) -> PredictResponse[TypeDocument]:
+        if get_type_var_name(document_class) != doc_config.document_class.__name__:
+            raise RuntimeError("Document class mismatch!")
+
         response = doc_config.document_class.request(
             doc_config.endpoints,
             self.input_doc,
@@ -106,7 +126,7 @@ def _make_request(
                 "API %s HTTP error: %s"
                 % (response.status_code, json.dumps(dict_response))
             )
-        return PredictResponse(
+        return PredictResponse[TypeDocument](
             http_response=dict_response,
             doc_config=doc_config,
             input_source=self.input_doc,
@@ -143,27 +163,27 @@ def __init__(self, api_key: str = "", raise_on_error: bool = True):
 
     def _init_default_endpoints(self) -> None:
         self._doc_configs = {
-            (OTS_OWNER, "invoice"): DocumentConfig(
-                document_type="invoice",
-                constructor=Invoice,
+            (OTS_OWNER, InvoiceV3.__name__): DocumentConfig(
+                document_type="invoice_v3",
+                document_class=InvoiceV3,
                 endpoints=[
                     StandardEndpoint(
                         url_name="invoices", version="3", api_key=self.api_key
                     )
                 ],
             ),
-            (OTS_OWNER, "receipt"): DocumentConfig(
-                document_type="receipt",
-                constructor=Receipt,
+            (OTS_OWNER, ReceiptV3.__name__): DocumentConfig(
+                document_type="receipt_v3",
+                document_class=ReceiptV3,
                 endpoints=[
                     StandardEndpoint(
                         url_name="expense_receipts", version="3", api_key=self.api_key
                     )
                 ],
             ),
-            (OTS_OWNER, "financial_doc"): DocumentConfig(
+            (OTS_OWNER, FinancialDocument.__name__): DocumentConfig(
                 document_type="financial_doc",
-                constructor=FinancialDocument,
+                document_class=FinancialDocument,
                 endpoints=[
                     StandardEndpoint(
                         url_name="invoices", version="3", api_key=self.api_key
@@ -173,18 +193,18 @@ def _init_default_endpoints(self) -> None:
                     ),
                 ],
             ),
-            (OTS_OWNER, "passport"): DocumentConfig(
-                document_type="passport",
-                constructor=Passport,
+            (OTS_OWNER, PassportV1.__name__): DocumentConfig(
+                document_type="passport_v1",
+                document_class=PassportV1,
                 endpoints=[
                     StandardEndpoint(
                         url_name="passport", version="1", api_key=self.api_key
                     )
                 ],
             ),
-            (OTS_OWNER, "bank_check"): DocumentConfig(
-                document_type="bank_check",
-                constructor=BankCheck,
+            (OTS_OWNER, BankCheckV1.__name__): DocumentConfig(
+                document_type="bank_check_v1",
+                document_class=BankCheckV1,
                 endpoints=[
                     StandardEndpoint(
                         url_name="bank_check", version="1", api_key=self.api_key
@@ -198,18 +218,21 @@ def add_endpoint(
         account_name: str,
         endpoint_name: str,
         version: str = "1",
+        document_class: Type[Document] = CustomV1,
     ) -> "Client":
         """
         Add a custom endpoint, created using the Mindee API Builder.
 
         :param endpoint_name: The "API name" field in the "Settings" page of the API Builder
         :param account_name: Your organization's username on the API Builder
         :param version: If set, locks the version of the model to use.
-                        If not set, use the latest version of the model.
+            If not set, use the latest version of the model.
+        :param document_class: A document class in which the response will be extracted.
+            Must inherit from ``mindee.documents.base.Document``.
         """
         self._doc_configs[(account_name, endpoint_name)] = DocumentConfig(
             document_type=endpoint_name,
-            constructor=CustomDocument,
+            document_class=document_class,
             endpoints=[
                 CustomEndpoint(
                     owner=account_name,

mindee/client.py

@@ -0,0 +1,6 @@
+from mindee.documents import us
+from mindee.documents.custom.custom_v1 import TypeCustomV1
+from mindee.documents.financial_document import TypeFinancialDocument
+from mindee.documents.invoice.invoice_v3 import TypeInvoiceV3
+from mindee.documents.passport.passport_v1 import TypePassportV1
+from mindee.documents.receipt.receipt_v3 import TypeReceiptV3