✨ add a specific class for classifications, which are never None (#134)

Author: ianardee

mindee/documents/custom/custom_v1.py

@@ -1,7 +1,7 @@
 from typing import Dict, Optional, TypeVar
 
 from mindee.documents.base import Document, TypeApiPrediction, clean_out_string
-from mindee.fields.api_builder import ClassificationField, ListField
+from mindee.documents.custom.custom_v1_fields import ClassificationField, ListField
 
 
 class CustomV1(Document):

mindee/fields/api_builder.py renamed to mindee/documents/custom/custom_v1_fields.py

@@ -3,6 +3,7 @@
 from mindee.documents.base import Document, TypeApiPrediction, clean_out_string
 from mindee.documents.invoice.line_item_v4 import InvoiceLineItemV4
 from mindee.fields.amount import AmountField
+from mindee.fields.classification import ClassificationField
 from mindee.fields.company_registration import CompanyRegistrationField
 from mindee.fields.date import DateField
 from mindee.fields.locale import LocaleField
@@ -50,11 +51,11 @@ class FinancialDocumentV1(Document):
     """Total amount of tip and gratuity."""
     time: TextField
     """Time as seen on the receipt in HH:MM format."""
-    document_type: TextField
+    document_type: ClassificationField
     """A classification field, among predefined classes."""
-    category: TextField
+    category: ClassificationField
     """The invoice or receipt category among predefined classes."""
-    subcategory: TextField
+    subcategory: ClassificationField
     """The invoice or receipt sub-category among predefined classes."""
 
     def __init__(
@@ -134,9 +135,13 @@ def _build_from_api_prediction(
         self.total_tax = AmountField(api_prediction["total_tax"], page_n=page_n)
         self.tip = AmountField(api_prediction["tip"], page_n=page_n)
         self.time = TextField(api_prediction["time"], page_n=page_n)
-        self.document_type = TextField(api_prediction["document_type"], page_n=page_n)
-        self.category = TextField(api_prediction["category"], page_n=page_n)
-        self.subcategory = TextField(api_prediction["subcategory"], page_n=page_n)
+        self.document_type = ClassificationField(
+            api_prediction["document_type"], page_n=page_n
+        )
+        self.category = ClassificationField(api_prediction["category"], page_n=page_n)
+        self.subcategory = ClassificationField(
+            api_prediction["subcategory"], page_n=page_n
+        )
 
     def __str__(self) -> str:
         supplier_company_registrations = "; ".join(

mindee/documents/financial/financial_document_v1.py

@@ -3,6 +3,7 @@
 from mindee.documents.base import Document, TypeApiPrediction, clean_out_string
 from mindee.documents.invoice import checks, reconstruct
 from mindee.fields.amount import AmountField
+from mindee.fields.classification import ClassificationField
 from mindee.fields.company_registration import CompanyRegistrationField
 from mindee.fields.date import DateField
 from mindee.fields.locale import LocaleField
@@ -14,6 +15,8 @@
 class InvoiceV3(Document):
     locale: LocaleField
     """locale information"""
+    document_type: ClassificationField
+    """Whether the document is an INVOICE or a CREDIT NOTE."""
     total_amount: AmountField
     """Total including taxes. Same as ``total_incl``."""
     total_net: AmountField
@@ -75,6 +78,9 @@ def _build_from_api_prediction(
         :param api_prediction: Raw prediction from HTTP response
         :param page_n: Page number for multi pages pdf input
         """
+        self.document_type = ClassificationField(
+            api_prediction["document_type"], page_n=page_n
+        )
         self.company_number = [
             CompanyRegistrationField(field_dict, page_n=page_n)
             for field_dict in api_prediction["company_registration"]

mindee/documents/invoice/invoice_v3.py

@@ -4,6 +4,7 @@
 from mindee.documents.invoice import checks, reconstruct
 from mindee.documents.invoice.line_item_v4 import InvoiceLineItemV4
 from mindee.fields.amount import AmountField
+from mindee.fields.classification import ClassificationField
 from mindee.fields.company_registration import CompanyRegistrationField
 from mindee.fields.date import DateField
 from mindee.fields.locale import LocaleField
@@ -15,6 +16,8 @@
 class InvoiceV4(Document):
     locale: LocaleField
     """locale information"""
+    document_type: ClassificationField
+    """Whether the document is an INVOICE or a CREDIT NOTE."""
     total_amount: AmountField
     """Total including taxes"""
     total_net: AmountField
@@ -80,6 +83,9 @@ def _build_from_api_prediction(
         :param api_prediction: Raw prediction from HTTP response
         :param page_n: Page number for multi pages pdf input
         """
+        self.document_type = ClassificationField(
+            api_prediction["document_type"], page_n=page_n
+        )
         self.supplier_company_registrations = [
             CompanyRegistrationField(field_dict, page_n=page_n)
             for field_dict in api_prediction["supplier_company_registrations"]

mindee/documents/invoice/invoice_v4.py

@@ -3,6 +3,7 @@
 from mindee.documents.base import Document, TypeApiPrediction, clean_out_string
 from mindee.fields.amount import AmountField
 from mindee.fields.base import field_array_confidence, field_array_sum
+from mindee.fields.classification import ClassificationField
 from mindee.fields.date import DateField
 from mindee.fields.locale import LocaleField
 from mindee.fields.tax import TaxField
@@ -18,7 +19,7 @@ class ReceiptV3(Document):
     """Date the receipt was issued"""
     time: TextField
     """Time the receipt was issued"""
-    category: TextField
+    category: ClassificationField
     """Service category"""
     merchant_name: TextField
     """Merchant's name"""
@@ -81,7 +82,7 @@ def _build_from_api_prediction(
         self.locale = LocaleField(api_prediction["locale"], page_n=page_n)
         self.total_incl = AmountField(api_prediction["total_incl"], page_n=page_n)
         self.date = DateField(api_prediction["date"], page_n=page_n)
-        self.category = TextField(api_prediction["category"], page_n=page_n)
+        self.category = ClassificationField(api_prediction["category"], page_n=page_n)
         self.merchant_name = TextField(
             api_prediction["supplier"], value_key="value", page_n=page_n
         )

mindee/documents/receipt/receipt_v3.py

@@ -2,6 +2,7 @@
 
 from mindee.documents.base import Document, TypeApiPrediction, clean_out_string
 from mindee.fields.amount import AmountField
+from mindee.fields.classification import ClassificationField
 from mindee.fields.date import DateField
 from mindee.fields.locale import LocaleField
 from mindee.fields.tax import TaxField
@@ -17,11 +18,11 @@ class ReceiptV4(Document):
     """Date the receipt was issued"""
     time: TextField
     """Time the receipt was issued, in HH: MM format."""
-    category: TextField
+    category: ClassificationField
     """The type, or service category, of the purchase."""
-    subcategory: TextField
+    subcategory: ClassificationField
     """The receipt sub category among predefined classes."""
-    document_type: TextField
+    document_type: ClassificationField
     """Whether the document is an expense receipt or a credit card receipt."""
     supplier: TextField
     """The merchant, or supplier, as found on the receipt."""
@@ -70,9 +71,13 @@ def _build_from_api_prediction(
         self.total_tax = AmountField(api_prediction["total_tax"], page_n=page_n)
         self.tip = AmountField(api_prediction["tip"], page_n=page_n)
         self.date = DateField(api_prediction["date"], page_n=page_n)
-        self.category = TextField(api_prediction["category"], page_n=page_n)
-        self.subcategory = TextField(api_prediction["subcategory"], page_n=page_n)
-        self.document_type = TextField(api_prediction["document_type"], page_n=page_n)
+        self.category = ClassificationField(api_prediction["category"], page_n=page_n)
+        self.subcategory = ClassificationField(
+            api_prediction["subcategory"], page_n=page_n
+        )
+        self.document_type = ClassificationField(
+            api_prediction["document_type"], page_n=page_n
+        )
         self.supplier = TextField(
             api_prediction["supplier"], value_key="value", page_n=page_n
         )

mindee/documents/receipt/receipt_v4.py

@@ -3,6 +3,7 @@
 from mindee.documents.base import Document, TypeApiPrediction, clean_out_string
 from mindee.documents.receipt.line_item_v5 import ReceiptV5LineItem
 from mindee.fields.amount import AmountField
+from mindee.fields.classification import ClassificationField
 from mindee.fields.company_registration import CompanyRegistrationField
 from mindee.fields.date import DateField
 from mindee.fields.locale import LocaleField
@@ -19,11 +20,11 @@ class ReceiptV5(Document):
     """The date the purchase was made."""
     time: TextField
     """Time of purchase with 24 hours formatting (HH:MM)."""
-    category: TextField
+    category: ClassificationField
     """The receipt category among predefined classes."""
-    subcategory: TextField
+    subcategory: ClassificationField
     """The receipt sub category among predefined classes for transport and food."""
-    document_type: TextField
+    document_type: ClassificationField
     """Whether the document is an expense receipt or a credit card receipt."""
     supplier_name: TextField
     """The name of the supplier or merchant."""
@@ -80,9 +81,13 @@ def _build_from_api_prediction(
         self.total_tax = AmountField(api_prediction["total_tax"], page_n=page_n)
         self.tip = AmountField(api_prediction["tip"], page_n=page_n)
         self.date = DateField(api_prediction["date"], page_n=page_n)
-        self.category = TextField(api_prediction["category"], page_n=page_n)
-        self.subcategory = TextField(api_prediction["subcategory"], page_n=page_n)
-        self.document_type = TextField(api_prediction["document_type"], page_n=page_n)
+        self.category = ClassificationField(api_prediction["category"], page_n=page_n)
+        self.subcategory = ClassificationField(
+            api_prediction["subcategory"], page_n=page_n
+        )
+        self.document_type = ClassificationField(
+            api_prediction["document_type"], page_n=page_n
+        )
         self.supplier_name = TextField(
             api_prediction["supplier_name"], value_key="value", page_n=page_n
         )

mindee/documents/receipt/receipt_v5.py

@@ -0,0 +1,32 @@
+from typing import Optional
+
+from mindee.fields.base import BaseField, TypePrediction
+
+
+class ClassificationField(BaseField):
+    """Represents a classifier value."""
+
+    value: str
+    """The value as a string."""
+
+    def __init__(
+        self,
+        prediction: TypePrediction,
+        value_key: str = "value",
+        reconstructed: bool = False,
+        page_n: Optional[int] = None,
+    ):
+        """
+        Text field object.
+
+        :param prediction: Amount prediction object from HTTP response
+        :param value_key: Key to use in the amount_prediction dict
+        :param reconstructed: Bool for reconstructed object (not extracted in the API)
+        :param page_n: Page number for multi-page document
+        """
+        super().__init__(
+            prediction,
+            value_key=value_key,
+            reconstructed=reconstructed,
+            page_n=page_n,
+        )