Update tests to use is_extracted

Author: qued

test_unstructured_inference/inference/test_layout.py

@@ -34,7 +34,7 @@ def mock_initial_layout():
         6,
         8,
         text="A very repetitive narrative. " * 10,
-        text_source="Mock",
+        is_extracted=True,
     )
 
     title_block = EmbeddedTextRegion.from_coords(
@@ -43,7 +43,7 @@ def mock_initial_layout():
         3,
         4,
         text="A Catchy Title",
-        text_source="Mock",
+        is_extracted=True,
     )
 
     return [text_block, title_block]

test_unstructured_inference/inference/test_layout_element.py

@@ -1,12 +1,12 @@
 from unstructured_inference.inference.layoutelement import LayoutElement, TextRegion
-from unstructured_inference.constants import Source, TextSource
+from unstructured_inference.constants import Source
 
 
 def test_layout_element_to_dict(mock_layout_element):
     expected = {
         "coordinates": ((100, 100), (100, 300), (300, 300), (300, 100)),
         "text": "Sample text",
-        "text_source": None,
+        "is_extracted": None,
         "type": "Text",
         "prob": None,
         "source": None,
@@ -28,7 +28,7 @@ def test_layoutelement_inheritance_works_correctly():
 
     # Create a TextRegion with both source and text_source
     region = TextRegion.from_coords(
-        0, 0, 10, 10, text="test", source=Source.YOLOX, text_source=TextSource.OCR
+        0, 0, 10, 10, text="test", source=Source.YOLOX, is_extracted=True
     )
 
     # Convert to LayoutElement
@@ -37,13 +37,13 @@ def test_layoutelement_inheritance_works_correctly():
     # Check that both properties are preserved
     assert element.source == Source.YOLOX, "LayoutElement should inherit source from TextRegion"
     assert (
-        element.text_source == TextSource.OCR
-    ), "LayoutElement should inherit text_source from TextRegion"
+        element.is_extracted
+    ), "LayoutElement should inherit is_extracted from TextRegion"
 
     # Check that to_dict() works correctly
     d = element.to_dict()
     assert d["source"] == Source.YOLOX
-    assert d["text_source"] == TextSource.OCR
+    assert d["is_extracted"]
 
     # Check that we can set source directly on LayoutElement
     element.source = Source.DETECTRON2_ONNX

test_unstructured_inference/test_elements.py

@@ -5,7 +5,7 @@
 import numpy as np
 import pytest
 
-from unstructured_inference.constants import Source, TextSource
+from unstructured_inference.constants import Source
 from unstructured_inference.inference import elements
 from unstructured_inference.inference.elements import (
     Rectangle,
@@ -410,8 +410,8 @@ def test_layoutelements_from_list_no_elements():
 
 def test_textregions_from_list_no_elements():
     back = TextRegions.from_list(regions=[])
-    assert back.text_sources.size == 0
-    assert back.text_source is None
+    assert back.is_extracted_array.size == 0
+    assert back.is_extracted is None
     assert back.element_coords.size == 0
 
 
@@ -456,8 +456,8 @@ def test_layoutelements_concatenate():
                 ]
             ),
             texts=np.array(["0", "1", "2", "3", "4"]),
-            text_sources=np.array([TextSource.OCR] * 5),
-            text_source=TextSource.OCR,
+            is_extracted_array=np.array([True] * 5),
+            is_extracted=True,
         ),
         LayoutElements(
             element_coords=np.array(
@@ -472,8 +472,8 @@ def test_layoutelements_concatenate():
             texts=np.array(["0", "1", "2", "3", "4"]),
             sources=np.array([Source.YOLOX] * 5),
             source=Source.YOLOX,
-            text_sources=np.array([TextSource.OCR] * 5),
-            text_source=TextSource.OCR,
+            is_extracted_array=np.array([] * 5),
+            is_extracted=True,
             element_probs=np.array([0.0, 0.1, 0.2, 0.3, 0.4]),
         ),
     ],
@@ -496,7 +496,7 @@ def test_textregions_from_list_collects_sources():
 
     regions = [
         TextRegion.from_coords(
-            0, 0, 10, 10, text="first", source=Source.YOLOX, text_source=TextSource.OCR
+            0, 0, 10, 10, text="first", source=Source.YOLOX, is_extracted=True
         ),
         TextRegion.from_coords(
             10,
@@ -505,7 +505,7 @@ def test_textregions_from_list_collects_sources():
             20,
             text="second",
             source=Source.DETECTRON2_ONNX,
-            text_source=TextSource.EXTRACTED,
+            is_extracted=True,
         ),
     ]
 
@@ -532,7 +532,7 @@ def test_textregions_iter_elements_preserves_source():
 
     regions = [
         TextRegion.from_coords(
-            0, 0, 10, 10, text="first", source=Source.YOLOX, text_source=TextSource.OCR
+            0, 0, 10, 10, text="first", source=Source.YOLOX, is_extracted=True
         ),
     ]
     text_regions = TextRegions.from_list(regions)
@@ -549,7 +549,7 @@ def test_textregions_slice_preserves_sources():
 
     regions = [
         TextRegion.from_coords(
-            0, 0, 10, 10, text="first", source=Source.YOLOX, text_source=TextSource.OCR
+            0, 0, 10, 10, text="first", source=Source.YOLOX, is_extracted=True
         ),
         TextRegion.from_coords(
             10,
@@ -558,7 +558,7 @@ def test_textregions_slice_preserves_sources():
             20,
             text="second",
             source=Source.DETECTRON2_ONNX,
-            text_source=TextSource.EXTRACTED,
+            is_extracted=True,
         ),
     ]
     text_regions = TextRegions.from_list(regions)
@@ -587,8 +587,8 @@ def test_textregions_from_coords_accepts_source():
     """Test that TextRegion.from_coords() accepts source parameter"""
     # This should fail because from_coords() doesn't accept source parameter
     region = TextRegion.from_coords(
-        0, 0, 10, 10, text="test", source=Source.YOLOX, text_source=TextSource.OCR
+        0, 0, 10, 10, text="test", source=Source.YOLOX, is_extracted=True
     )
 
     assert region.source == Source.YOLOX
-    assert region.text_source == TextSource.OCR
+    assert region.is_extracted