add test

Author: valerie-fauconmorin-kensho

kensho_kenverters/CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## v1.2.3
+
+Remove table validation for rows and columns to not fail downstream of Extract model failures
+
 ## v1.2.2
 
 Use underscores instead of hyphens for setuptools 78

kensho_kenverters/tables_utils.py

@@ -23,7 +23,7 @@ def _create_empty_annotation(row: int, col: int) -> AnnotationModel:
 
 
 def _validate_annotations(
-    duplicated_annotations: list[AnnotationModel],
+    duplicated_annotations: list[AnnotationModel], max_row: int, max_col: int
 ) -> list[AnnotationModel]:
     """Validate duplicated annotations."""
     # Check all spans are 1 (annotations are duplicated)
@@ -36,19 +36,10 @@ def _validate_annotations(
     if len(set(all_indices)) != len(all_indices):
         raise ValueError("Overlapping indices in table.")
 
-    # If any empty rows or columns, add empty annotations
-    existing_rows = set(index[0] for index in all_indices)
-    existing_columns = set(index[1] for index in all_indices)
-    full_rows = set(range(max(existing_rows, default=0) + 1))
-    full_columns = set(range(max(existing_columns, default=0) + 1))
-
-    for row in full_rows - existing_rows:
-        for col in full_columns:
-            duplicated_annotations.append(_create_empty_annotation(row, col))
-
-    for col in full_columns - existing_columns:
-        for row in full_rows:
-            duplicated_annotations.append(_create_empty_annotation(row, col))
+    for row in range(max_row + 1):
+        for col in range(max_col + 1):
+            if (row, col) not in all_indices:
+                duplicated_annotations.append(_create_empty_annotation(row, col))
 
     return duplicated_annotations
 
@@ -70,6 +61,8 @@ def duplicate_spanning_annotations(
         duplicated annotations. Duplicated annotations must all have span (1, 1).
     """
     duplicated_annotations = []
+    max_row = 0
+    max_col = 0
     for annotation in annotations:
         data = annotation.data
         row_span, col_span = data.span
@@ -94,8 +87,10 @@ def duplicate_spanning_annotations(
                     locations=annotation.locations,
                 )
                 duplicated_annotations.append(new_annotation)
+                max_row = max(max_row, row_index + row_span_index)
+                max_col = max(max_col, col_index + col_span_index)
 
-    return _validate_annotations(duplicated_annotations)
+    return _validate_annotations(duplicated_annotations, max_row, max_col)
 
 
 def get_table_shape(

kensho_kenverters/tests/test_tables_utils.py

@@ -57,7 +57,7 @@ def test_convert_tpdle_to_pd_df(self) -> None:
         )
         self.assertEqual(expected_df.to_csv(), converted_df.to_csv())
 
-    def test_duplicate_spanning_annotations(self) -> None:
+    def test_duplicate_spanning_annotations_1_spans(self) -> None:
         # Test no duplication when all spans are 1
         duplicated = duplicate_spanning_annotations(
             self.parsed_serialized_document.annotations
@@ -417,6 +417,7 @@ def test_duplicate_spanning_annotations(self) -> None:
 
         self.assertEqual(duplicated, expected_duplicated)
 
+    def test_duplicate_spanning_annotations_greater_1_spans(self) -> None:
         # Test properly duplicated when cells have a span > 1
         annotations = self.parsed_serialized_document.annotations
         annotations.pop(-1)
@@ -775,3 +776,53 @@ def test_duplicate_spanning_annotations(self) -> None:
         ]
         duplicated = duplicate_spanning_annotations(annotations)
         self.assertEqual(expected_duplicated, duplicated)
+
+    def test_duplicate_spanning_annotations_missing_cells(self) -> None:
+        # Test missing cells
+        annotations = [
+            AnnotationModel(
+                content_uids=["7"],
+                data=AnnotationDataModel(index=(0, 0), span=(1, 1)),
+                type="table_structure",
+                locations=[
+                    LocationModel(
+                        height=0.01188,
+                        width=0.22128,
+                        x=0.16008,
+                        y=0.40464,
+                        page_number=0,
+                    )
+                ],
+            ),
+            AnnotationModel(
+                content_uids=["8"],
+                data=AnnotationDataModel(index=(0, 1), span=(1, 1)),
+                type="table_structure",
+                locations=[
+                    LocationModel(
+                        height=0.01188,
+                        width=0.22128,
+                        x=0.16008,
+                        y=0.80464,
+                        page_number=0,
+                    )
+                ],
+            ),
+            AnnotationModel(
+                content_uids=["9"],
+                data=AnnotationDataModel(index=(3, 0), span=(1, 1)),
+                type="table_structure",
+                locations=[
+                    LocationModel(
+                        height=0.01188,
+                        width=0.22128,
+                        x=0.7008,
+                        y=0.40464,
+                        page_number=0,
+                    )
+                ],
+            ),
+        ]
+
+        duplicated = duplicate_spanning_annotations(annotations)
+        self.assertEqual(len(duplicated), 8)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "kensho_kenverters"
-version = "1.2.2"
+version = "1.2.3"
 description = "Extract Output Translator Tools"
 readme = "README.md"
 authors = ["Valerie Faucon-Morin <valerie.fauconmorin@kensho.com>"]