🐛 page_n should always be set when available (#106)

Author: ianardee

mindee/fields/api_builder.py

@@ -37,7 +37,7 @@ class ListField:
     """Confidence score"""
     reconstructed: bool
     """Whether the field was reconstructed from other fields."""
-    page_n: Optional[int] = None
+    page_n: int
     """The document page on which the information was found."""
     values: List[ListFieldValue]
     """List of word values"""
@@ -50,7 +50,10 @@ def __init__(
     ):
         self.values = []
         self.reconstructed = reconstructed
-        self.page_n = page_n
+        if page_n is None:
+            self.page_n = prediction["page_id"]
+        else:
+            self.page_n = page_n
         self.confidence = prediction["confidence"]
 
         for value in prediction["values"]:

mindee/fields/base.py

@@ -49,7 +49,14 @@ def __init__(
         :param reconstructed: Bool for reconstructed object (not extracted in the API)
         :param page_n: Page number for multi-page PDF
         """
-        self.page_n = page_n
+        if page_n is None:
+            try:
+                self.page_n = prediction["page_id"]
+            except KeyError:
+                pass
+        else:
+            self.page_n = page_n
+
         self.reconstructed = reconstructed
 
         if value_key not in prediction or prediction[value_key] == "N/A":

tests/documents/test_custom_v1.py

@@ -57,6 +57,7 @@ def test_complete(custom_v1_doc_object):
             assert value.content != ""
             assert len(value.bounding_box) == 4
             assert value.confidence != 0.0
+            assert field.page_n == 0 or field.page_n == 1
     for field_name, field in custom_v1_doc_object.classifications.items():
         assert len(field_name) > 0
         assert isinstance(field, ClassificationField)
@@ -66,3 +67,7 @@ def test_complete(custom_v1_doc_object):
 
 def test_page_complete(custom_v1_page_object):
     assert custom_v1_page_object.orientation.value == 0
+    for field_name, field in custom_v1_page_object.fields.items():
+        assert isinstance(field, ListField)
+        assert len(field.contents_list) == len(field.values)
+        assert field.page_n == 0

tests/documents/test_invoice_v3.py

@@ -20,14 +20,21 @@ def invoice_v3_doc_object_empty(invoice_pred):
     return InvoiceV3(api_prediction=invoice_pred)
 
 
+@pytest.fixture
+def invoice_v3_page_object():
+    json_data = json.load(open(FILE_PATH_INVOICE_V3_COMPLETE))
+    return InvoiceV3(
+        api_prediction=json_data["document"]["inference"]["pages"][0], page_n=0
+    )
+
+
 @pytest.fixture
 def invoice_pred():
     json_data = json.load(open(FILE_PATH_INVOICE_V3_EMPTY))
     return json_data["document"]["inference"]["pages"][0]
 
 
-# Technical tests
-def test_constructor(invoice_v3_doc_object):
+def test_doc_constructor(invoice_v3_doc_object):
     assert invoice_v3_doc_object.invoice_date.value == "2020-02-17"
     assert invoice_v3_doc_object.checklist["taxes_match_total_incl"] is True
     assert invoice_v3_doc_object.checklist["taxes_match_total_excl"] is True
@@ -44,6 +51,14 @@ def test_constructor(invoice_v3_doc_object):
     assert str(invoice_v3_doc_object) == doc_str
 
 
+def test_page_constructor(invoice_v3_page_object):
+    doc_str = open(f"{INVOICE_DATA_DIR}/response_v3/page0_to_string.txt").read().strip()
+    assert invoice_v3_page_object.orientation.value == 0
+    assert invoice_v3_page_object.invoice_number.page_n == 0
+    assert str(invoice_v3_page_object) == doc_str
+    assert len(invoice_v3_page_object.cropper) == 0
+
+
 def test_all_na(invoice_v3_doc_object_empty):
     assert invoice_v3_doc_object_empty.locale.value is None
     assert invoice_v3_doc_object_empty.total_incl.value is None

tests/documents/test_passport_v1.py

@@ -11,13 +11,21 @@
 @pytest.fixture
 def passport_v1_doc_object():
     json_data = json.load(open(FILE_PATH_PASSPORT_V1_COMPLETE))
-    return PassportV1(api_prediction=json_data["document"]["inference"]["pages"][0])
+    return PassportV1(api_prediction=json_data["document"]["inference"], page_n=None)
 
 
 @pytest.fixture
 def passport_v1_doc_object_empty():
     json_data = json.load(open(f"{PASSPORT_DATA_DIR}/response_v1/empty.json"))
-    return PassportV1(api_prediction=json_data["document"]["inference"]["pages"][0])
+    return PassportV1(api_prediction=json_data["document"]["inference"], page_n=None)
+
+
+@pytest.fixture
+def passport_v1_page_object():
+    json_data = json.load(open(FILE_PATH_PASSPORT_V1_COMPLETE))
+    return PassportV1(
+        api_prediction=json_data["document"]["inference"]["pages"][0], page_n=0
+    )
 
 
 def test_constructor(passport_v1_doc_object):
@@ -26,9 +34,20 @@ def test_constructor(passport_v1_doc_object):
     doc_str = (
         open(f"{PASSPORT_DATA_DIR}/response_v1/page0_to_string.txt").read().strip()
     )
+    assert passport_v1_doc_object.birth_date.page_n == 0
     assert str(passport_v1_doc_object) == doc_str
 
 
+def test_page_constructor(passport_v1_page_object):
+    doc_str = (
+        open(f"{PASSPORT_DATA_DIR}/response_v1/page0_to_string.txt").read().strip()
+    )
+    assert passport_v1_page_object.orientation.value == 0
+    assert passport_v1_page_object.birth_date.page_n == 0
+    assert str(passport_v1_page_object) == doc_str
+    assert len(passport_v1_page_object.cropper) == 0
+
+
 def test_all_na(passport_v1_doc_object_empty):
     assert passport_v1_doc_object_empty.mrz.value is None
     assert passport_v1_doc_object_empty.country.value is None

tests/documents/test_receipt_v3.py

@@ -27,12 +27,12 @@ def receipt_pred():
     return json_data["document"]["inference"]["pages"][0]
 
 
-# Technical tests
-def test_constructor(receipt_v3_doc_object):
+def test_doc_constructor(receipt_v3_doc_object):
     assert receipt_v3_doc_object.date.value == "2016-02-26"
     assert receipt_v3_doc_object.total_tax.value == 1.7
     assert receipt_v3_doc_object.checklist["taxes_match_total_incl"] is True
     doc_str = open(f"{RECEIPT_DATA_DIR}/response_v3/doc_to_string.txt").read().strip()
+    assert receipt_v3_doc_object.date.page_n == 0
     assert str(receipt_v3_doc_object) == doc_str
 
 

tests/documents/test_receipt_v4.py

@@ -34,6 +34,7 @@ def test_doc_constructor(receipt_v4_doc_object):
     assert receipt_v4_doc_object.total_tax.value == 3.34
     doc_str = open(f"{RECEIPT_DATA_DIR}/response_v4/doc_to_string.txt").read().strip()
     assert receipt_v4_doc_object.orientation is None
+    assert receipt_v4_doc_object.date.page_n == 0
     assert str(receipt_v4_doc_object) == doc_str
 
 
@@ -42,6 +43,7 @@ def test_page_constructor(receipt_v4_page_object):
     assert receipt_v4_page_object.total_tax.value == 3.34
     doc_str = open(f"{RECEIPT_DATA_DIR}/response_v4/page0_to_string.txt").read().strip()
     assert receipt_v4_page_object.orientation.value == 0
+    assert receipt_v4_page_object.date.page_n == 0
     assert str(receipt_v4_page_object) == doc_str
     assert len(receipt_v4_page_object.cropper) == 0