修改为东南天坐标系
This commit is contained in:
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,468 @@
|
||||
import pathlib
|
||||
from multiprocessing import Pool
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
from unstructured_inference.constants import IsExtracted
|
||||
from unstructured_inference.inference import layout
|
||||
from unstructured_inference.inference.elements import TextRegion
|
||||
from unstructured_inference.inference.layoutelement import LayoutElement
|
||||
|
||||
from test_unstructured.unit_utils import example_doc_path
|
||||
from unstructured.documents.coordinates import PixelSpace
|
||||
from unstructured.documents.elements import (
|
||||
TYPE_TO_TEXT_ELEMENT_MAP,
|
||||
CheckBox,
|
||||
CoordinatesMetadata,
|
||||
ElementType,
|
||||
FigureCaption,
|
||||
Header,
|
||||
ListItem,
|
||||
NarrativeText,
|
||||
Text,
|
||||
Title,
|
||||
)
|
||||
from unstructured.documents.elements import (
|
||||
Image as ImageElement,
|
||||
)
|
||||
from unstructured.partition.common import common
|
||||
|
||||
|
||||
class MockPageLayout(layout.PageLayout):
|
||||
def __init__(self, number: int, image: Image.Image):
|
||||
self.number = number
|
||||
self.image = image
|
||||
|
||||
@property
|
||||
def elements(self):
|
||||
return [
|
||||
LayoutElement(
|
||||
type="Headline",
|
||||
text="Charlie Brown and the Great Pumpkin",
|
||||
bbox=None,
|
||||
),
|
||||
LayoutElement(
|
||||
type="Subheadline",
|
||||
text="The Beginning",
|
||||
bbox=None,
|
||||
),
|
||||
LayoutElement(
|
||||
type="Text",
|
||||
text="This time Charlie Brown had it really tricky...",
|
||||
bbox=None,
|
||||
),
|
||||
LayoutElement(
|
||||
type="Title",
|
||||
text="Another book title in the same page",
|
||||
bbox=None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class MockDocumentLayout(layout.DocumentLayout):
|
||||
@property
|
||||
def pages(self):
|
||||
return [
|
||||
MockPageLayout(number=1, image=Image.new("1", (1, 1))),
|
||||
]
|
||||
|
||||
|
||||
def test_normalize_layout_element_dict():
|
||||
layout_element = {
|
||||
"type": "Title",
|
||||
"coordinates": [[1, 2], [3, 4], [5, 6], [7, 8]],
|
||||
"coordinate_system": None,
|
||||
"text": "Some lovely text",
|
||||
}
|
||||
coordinate_system = PixelSpace(width=10, height=20)
|
||||
element = common.normalize_layout_element(
|
||||
layout_element,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
assert element == Title(
|
||||
text="Some lovely text",
|
||||
coordinates=[[1, 2], [3, 4], [5, 6], [7, 8]],
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
|
||||
|
||||
def test_normalize_layout_element_dict_caption():
|
||||
layout_element = {
|
||||
"type": "Figure",
|
||||
"coordinates": ((1, 2), (3, 4), (5, 6), (7, 8)),
|
||||
"text": "Some lovely text",
|
||||
}
|
||||
coordinate_system = PixelSpace(width=10, height=20)
|
||||
element = common.normalize_layout_element(
|
||||
layout_element,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
assert element == ImageElement(
|
||||
text="Some lovely text",
|
||||
coordinates=((1, 2), (3, 4), (5, 6), (7, 8)),
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("element_type", "expected_type", "expected_depth"),
|
||||
[
|
||||
("Title", Title, None),
|
||||
("Headline", Title, 1),
|
||||
("Subheadline", Title, 2),
|
||||
("Header", Header, None),
|
||||
],
|
||||
)
|
||||
def test_normalize_layout_element_headline(element_type, expected_type, expected_depth):
|
||||
layout_element = {
|
||||
"type": element_type,
|
||||
"coordinates": [[1, 2], [3, 4], [5, 6], [7, 8]],
|
||||
"text": "Some lovely text",
|
||||
}
|
||||
coordinate_system = PixelSpace(width=10, height=20)
|
||||
element = common.normalize_layout_element(layout_element, coordinate_system=coordinate_system)
|
||||
assert element.metadata.category_depth == expected_depth
|
||||
assert isinstance(element, expected_type)
|
||||
|
||||
|
||||
def test_normalize_layout_element_dict_figure_caption():
|
||||
layout_element = {
|
||||
"type": "FigureCaption",
|
||||
"coordinates": [[1, 2], [3, 4], [5, 6], [7, 8]],
|
||||
"text": "Some lovely text",
|
||||
}
|
||||
coordinate_system = PixelSpace(width=10, height=20)
|
||||
element = common.normalize_layout_element(
|
||||
layout_element,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
assert element == FigureCaption(
|
||||
text="Some lovely text",
|
||||
coordinates=[[1, 2], [3, 4], [5, 6], [7, 8]],
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
|
||||
|
||||
def test_normalize_layout_element_dict_misc():
|
||||
layout_element = {
|
||||
"type": "Misc",
|
||||
"coordinates": [[1, 2], [3, 4], [5, 6], [7, 8]],
|
||||
"text": "Some lovely text",
|
||||
}
|
||||
coordinate_system = PixelSpace(width=10, height=20)
|
||||
element = common.normalize_layout_element(
|
||||
layout_element,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
assert element == Text(
|
||||
text="Some lovely text",
|
||||
coordinates=[[1, 2], [3, 4], [5, 6], [7, 8]],
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
|
||||
|
||||
def test_normalize_layout_element_layout_element():
|
||||
layout_element = LayoutElement.from_coords(
|
||||
type="Text",
|
||||
x1=1,
|
||||
y1=2,
|
||||
x2=3,
|
||||
y2=4,
|
||||
text="Some lovely text",
|
||||
)
|
||||
coordinate_system = PixelSpace(width=10, height=20)
|
||||
element = common.normalize_layout_element(
|
||||
layout_element,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
assert element == NarrativeText(
|
||||
text="Some lovely text",
|
||||
coordinates=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
|
||||
|
||||
def test_normalize_layout_element_layout_element_narrative_text():
|
||||
layout_element = LayoutElement.from_coords(
|
||||
type="NarrativeText",
|
||||
x1=1,
|
||||
y1=2,
|
||||
x2=3,
|
||||
y2=4,
|
||||
text="Some lovely text",
|
||||
)
|
||||
coordinate_system = PixelSpace(width=10, height=20)
|
||||
element = common.normalize_layout_element(
|
||||
layout_element,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
assert element == NarrativeText(
|
||||
text="Some lovely text",
|
||||
coordinates=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("element_type", "expected_element_class"),
|
||||
TYPE_TO_TEXT_ELEMENT_MAP.items(),
|
||||
)
|
||||
def test_normalize_layout_element_layout_element_maps_to_appropriate_text_element(
|
||||
element_type: str,
|
||||
expected_element_class: type[Text],
|
||||
):
|
||||
layout_element = LayoutElement.from_coords(
|
||||
type=element_type,
|
||||
x1=1,
|
||||
y1=2,
|
||||
x2=3,
|
||||
y2=4,
|
||||
text="Some lovely text",
|
||||
)
|
||||
coordinate_system = PixelSpace(width=10, height=20)
|
||||
element = common.normalize_layout_element(
|
||||
layout_element,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
assert element == expected_element_class(
|
||||
text="Some lovely text",
|
||||
coordinates=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("element_type", "expected_checked"),
|
||||
[
|
||||
(ElementType.CHECK_BOX_UNCHECKED, False),
|
||||
(ElementType.CHECK_BOX_CHECKED, True),
|
||||
(ElementType.RADIO_BUTTON_UNCHECKED, False),
|
||||
(ElementType.RADIO_BUTTON_CHECKED, True),
|
||||
(ElementType.CHECKED, True),
|
||||
(ElementType.UNCHECKED, False),
|
||||
],
|
||||
)
|
||||
def test_normalize_layout_element_checkable(element_type: str, expected_checked: bool):
|
||||
layout_element = LayoutElement.from_coords(
|
||||
type=element_type,
|
||||
x1=1,
|
||||
y1=2,
|
||||
x2=3,
|
||||
y2=4,
|
||||
text="",
|
||||
)
|
||||
coordinate_system = PixelSpace(width=10, height=20)
|
||||
element = common.normalize_layout_element(
|
||||
layout_element,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
assert isinstance(element, CheckBox)
|
||||
assert element == CheckBox(
|
||||
checked=expected_checked,
|
||||
coordinates=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
|
||||
|
||||
def test_normalize_layout_element_enumerated_list():
|
||||
layout_element = LayoutElement.from_coords(
|
||||
type="List",
|
||||
x1=1,
|
||||
y1=2,
|
||||
x2=3,
|
||||
y2=4,
|
||||
text="1. I'm so cool! 2. You're cool too. 3. We're all cool!",
|
||||
)
|
||||
coordinate_system = PixelSpace(width=10, height=20)
|
||||
elements = common.normalize_layout_element(
|
||||
layout_element,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
assert elements == [
|
||||
ListItem(
|
||||
text="I'm so cool!",
|
||||
coordinates=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
coordinate_system=coordinate_system,
|
||||
),
|
||||
ListItem(
|
||||
text="You're cool too.",
|
||||
coordinates=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
coordinate_system=coordinate_system,
|
||||
),
|
||||
ListItem(
|
||||
text="We're all cool!",
|
||||
coordinates=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
coordinate_system=coordinate_system,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_normalize_layout_element_bulleted_list():
|
||||
layout_element = LayoutElement.from_coords(
|
||||
type="List",
|
||||
x1=1,
|
||||
y1=2,
|
||||
x2=3,
|
||||
y2=4,
|
||||
text="* I'm so cool! * You're cool too. * We're all cool!",
|
||||
)
|
||||
coordinate_system = PixelSpace(width=10, height=20)
|
||||
elements = common.normalize_layout_element(
|
||||
layout_element,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
assert elements == [
|
||||
ListItem(
|
||||
text="I'm so cool!",
|
||||
coordinates=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
coordinate_system=coordinate_system,
|
||||
),
|
||||
ListItem(
|
||||
text="You're cool too.",
|
||||
coordinates=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
coordinate_system=coordinate_system,
|
||||
),
|
||||
ListItem(
|
||||
text="We're all cool!",
|
||||
coordinates=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
coordinate_system=coordinate_system,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class MockRunOutput:
|
||||
def __init__(self, returncode, stdout, stderr):
|
||||
self.returncode = returncode
|
||||
self.stdout = stdout
|
||||
self.stderr = stderr
|
||||
|
||||
|
||||
def test_convert_office_doc_captures_errors(monkeypatch, caplog):
|
||||
from unstructured.partition.common.common import subprocess
|
||||
|
||||
def mock_run(*args, **kwargs):
|
||||
return MockRunOutput(1, "an error occurred".encode(), "error details".encode())
|
||||
|
||||
monkeypatch.setattr(subprocess, "run", mock_run)
|
||||
common.convert_office_doc("no-real.docx", "fake-directory", target_format="docx")
|
||||
assert "soffice failed to convert to format docx with code 1" in caplog.text
|
||||
|
||||
|
||||
def test_convert_office_docs_avoids_concurrent_call_to_soffice():
|
||||
paths_to_save = [pathlib.Path(path) for path in ("/tmp/proc1", "/tmp/proc2", "/tmp/proc3")]
|
||||
for path in paths_to_save:
|
||||
path.mkdir(exist_ok=True)
|
||||
(path / "simple.docx").unlink(missing_ok=True)
|
||||
file_to_convert = example_doc_path("simple.doc")
|
||||
|
||||
with Pool(3) as pool:
|
||||
pool.starmap(common.convert_office_doc, [(file_to_convert, path) for path in paths_to_save])
|
||||
|
||||
assert np.sum([(path / "simple.docx").is_file() for path in paths_to_save]) == 3
|
||||
|
||||
|
||||
def test_convert_office_docs_respects_wait_timeout():
|
||||
paths_to_save = [
|
||||
pathlib.Path(path) for path in ("/tmp/wait/proc1", "/tmp/wait/proc2", "/tmp/wait/proc3")
|
||||
]
|
||||
for path in paths_to_save:
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
(path / "simple.docx").unlink(missing_ok=True)
|
||||
file_to_convert = example_doc_path("simple.doc")
|
||||
|
||||
with Pool(3) as pool:
|
||||
pool.starmap(
|
||||
common.convert_office_doc,
|
||||
# set timeout to wait for soffice to be available to 0 so only one process can convert
|
||||
# the doc file on the first try; then the catch all
|
||||
[(file_to_convert, path, "docx", None, 0) for path in paths_to_save],
|
||||
)
|
||||
|
||||
# because this test file is very small we could have occasions where two files are converted
|
||||
# when one of the processes spawned just a little
|
||||
assert np.sum([(path / "simple.docx").is_file() for path in paths_to_save]) < 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("<table><tbody><tr><td>👨\\U+1F3FB🔧</td></tr></tbody></table>", True),
|
||||
("<table><tbody><tr><td>Hello!</td></tr></tbody></table>", False),
|
||||
],
|
||||
)
|
||||
def test_contains_emoji(text, expected):
|
||||
assert common.contains_emoji(text) is expected
|
||||
|
||||
|
||||
def test_get_page_image_metadata_and_coordinate_system():
|
||||
doc = MockDocumentLayout()
|
||||
metadata = common.get_page_image_metadata(doc.pages[0])
|
||||
assert isinstance(metadata, dict)
|
||||
|
||||
|
||||
def test_ocr_data_to_elements():
|
||||
filename = example_doc_path("img/layout-parser-paper-fast.jpg")
|
||||
text_regions = [
|
||||
TextRegion.from_coords(
|
||||
163.0,
|
||||
115.0,
|
||||
452.0,
|
||||
129.0,
|
||||
text="LayoutParser: A Unified Toolkit for Deep",
|
||||
),
|
||||
TextRegion.from_coords(
|
||||
156.0,
|
||||
132.0,
|
||||
457.0,
|
||||
147.0,
|
||||
text="Learning Based Document Image Analysis",
|
||||
),
|
||||
]
|
||||
ocr_data = [
|
||||
LayoutElement(
|
||||
bbox=r.bbox,
|
||||
text=r.text,
|
||||
source=r.source,
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
)
|
||||
for r in text_regions
|
||||
]
|
||||
image = Image.open(filename)
|
||||
|
||||
elements = common.ocr_data_to_elements(
|
||||
ocr_data=ocr_data,
|
||||
image_size=image.size,
|
||||
)
|
||||
|
||||
assert len(ocr_data) == len(elements)
|
||||
assert {el.category for el in elements} == {ElementType.UNCATEGORIZED_TEXT}
|
||||
|
||||
# check coordinates metadata
|
||||
image_width, image_height = image.size
|
||||
coordinate_system = PixelSpace(width=image_width, height=image_height)
|
||||
for el, layout_el in zip(elements, ocr_data):
|
||||
assert el.metadata.coordinates == CoordinatesMetadata(
|
||||
points=layout_el.bbox.coordinates,
|
||||
system=coordinate_system,
|
||||
)
|
||||
|
||||
|
||||
def test_normalize_layout_element_layout_element_text_source_metadata():
|
||||
layout_element = LayoutElement.from_coords(
|
||||
type="NarrativeText",
|
||||
x1=1,
|
||||
y1=2,
|
||||
x2=3,
|
||||
y2=4,
|
||||
text="Some lovely text",
|
||||
is_extracted=IsExtracted.TRUE,
|
||||
)
|
||||
coordinate_system = PixelSpace(width=10, height=20)
|
||||
element = common.normalize_layout_element(
|
||||
layout_element,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
assert hasattr(element, "metadata")
|
||||
assert hasattr(element.metadata, "is_extracted")
|
||||
assert element.metadata.is_extracted == "true"
|
||||
@@ -0,0 +1,282 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
"""Unit-test suite for the `unstructured.partition.lang` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
|
||||
import pytest
|
||||
|
||||
from test_unstructured.unit_utils import LogCaptureFixture
|
||||
from unstructured.documents.elements import (
|
||||
NarrativeText,
|
||||
PageBreak,
|
||||
)
|
||||
from unstructured.partition.common.lang import (
|
||||
_clean_ocr_languages_arg,
|
||||
_convert_language_code_to_pytesseract_lang_code,
|
||||
apply_lang_metadata,
|
||||
check_language_args,
|
||||
detect_languages,
|
||||
prepare_languages_for_tesseract,
|
||||
tesseract_to_paddle_language,
|
||||
)
|
||||
|
||||
DIRECTORY = pathlib.Path(__file__).parent.resolve()
|
||||
EXAMPLE_DOCS_DIRECTORY = os.path.join(DIRECTORY, "..", "..", "example-docs")
|
||||
|
||||
|
||||
def test_prepare_languages_for_tesseract_with_one_language():
|
||||
languages = ["en"]
|
||||
assert prepare_languages_for_tesseract(languages) == "eng"
|
||||
|
||||
|
||||
def test_prepare_languages_for_tesseract_with_duplicated_languages():
|
||||
languages = ["en", "eng"]
|
||||
assert prepare_languages_for_tesseract(languages) == "eng"
|
||||
|
||||
|
||||
def test_prepare_languages_for_tesseract_special_case():
|
||||
languages = ["osd"]
|
||||
assert prepare_languages_for_tesseract(languages) == "osd"
|
||||
|
||||
languages = ["equ"]
|
||||
assert prepare_languages_for_tesseract(languages) == "equ"
|
||||
|
||||
|
||||
def test_prepare_languages_for_tesseract_removes_empty_inputs():
|
||||
languages = ["kbd", "es"]
|
||||
assert prepare_languages_for_tesseract(languages) == "spa+spa_old"
|
||||
|
||||
|
||||
def test_prepare_languages_for_tesseract_includes_variants():
|
||||
languages = ["chi"]
|
||||
assert prepare_languages_for_tesseract(languages) == "chi_sim+chi_sim_vert+chi_tra+chi_tra_vert"
|
||||
|
||||
|
||||
def test_prepare_languages_for_tesseract_with_multiple_languages():
|
||||
languages = ["ja", "afr", "en", "equ"]
|
||||
assert prepare_languages_for_tesseract(languages) == "jpn+jpn_vert+afr+eng+equ"
|
||||
|
||||
|
||||
def test_prepare_languages_for_tesseract_warns_nonstandard_language(caplog: LogCaptureFixture):
|
||||
languages = ["zzz", "chi"]
|
||||
assert prepare_languages_for_tesseract(languages) == "chi_sim+chi_sim_vert+chi_tra+chi_tra_vert"
|
||||
assert "not a valid standard language code" in caplog.text
|
||||
|
||||
|
||||
def test_prepare_languages_for_tesseract_warns_non_tesseract_language(caplog: LogCaptureFixture):
|
||||
languages = ["kbd", "eng"]
|
||||
assert prepare_languages_for_tesseract(languages) == "eng"
|
||||
assert "not a language supported by Tesseract" in caplog.text
|
||||
|
||||
|
||||
def test_prepare_languages_for_tesseract_None_languages():
|
||||
with pytest.raises(ValueError, match="`languages` can not be `None`"):
|
||||
languages = None
|
||||
prepare_languages_for_tesseract(languages)
|
||||
|
||||
|
||||
def test_prepare_languages_for_tesseract_no_valid_languages(caplog: LogCaptureFixture):
|
||||
languages = [""]
|
||||
assert prepare_languages_for_tesseract(languages) == "eng"
|
||||
assert "Failed to find any valid standard language code from languages" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tesseract_lang", "expected_lang"),
|
||||
[
|
||||
("eng", "en"),
|
||||
("chi_sim", "ch"),
|
||||
("chi_tra", "chinese_cht"),
|
||||
("deu", "german"),
|
||||
("jpn", "japan"),
|
||||
("kor", "korean"),
|
||||
],
|
||||
)
|
||||
def test_tesseract_to_paddle_language_valid_codes(tesseract_lang: str, expected_lang: str):
|
||||
assert expected_lang == tesseract_to_paddle_language(tesseract_lang)
|
||||
|
||||
|
||||
def test_tesseract_to_paddle_language_invalid_codes(caplog: LogCaptureFixture):
|
||||
tesseract_lang = "unsupported_lang"
|
||||
assert tesseract_to_paddle_language(tesseract_lang) == "en"
|
||||
assert "unsupported_lang is not a language code supported by PaddleOCR," in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tesseract_lang", "expected_lang"),
|
||||
[
|
||||
("ENG", "en"),
|
||||
("Fra", "fr"),
|
||||
("DEU", "german"),
|
||||
],
|
||||
)
|
||||
def test_tesseract_to_paddle_language_case_sensitivity(tesseract_lang: str, expected_lang: str):
|
||||
assert expected_lang == tesseract_to_paddle_language(tesseract_lang)
|
||||
|
||||
|
||||
def test_detect_languages_english_auto():
|
||||
text = "This is a short sentence."
|
||||
assert detect_languages(text) == ["eng"]
|
||||
|
||||
|
||||
def test_detect_languages_english_provided():
|
||||
text = "This is another short sentence."
|
||||
languages = ["en"]
|
||||
assert detect_languages(text, languages) == ["eng"]
|
||||
|
||||
|
||||
def test_detect_languages_korean_auto():
|
||||
text = "안녕하세요"
|
||||
assert detect_languages(text) == ["kor"]
|
||||
|
||||
|
||||
def test_detect_languages_gets_multiple_languages():
|
||||
text = "My lubimy mleko i chleb."
|
||||
assert detect_languages(text) == ["ces", "pol", "slk"]
|
||||
|
||||
|
||||
def test_detect_languages_warns_for_auto_and_other_input(caplog: LogCaptureFixture):
|
||||
text = "This is another short sentence."
|
||||
languages = ["en", "auto", "rus"]
|
||||
assert detect_languages(text, languages) == ["eng"]
|
||||
assert "rest of the inputted languages will be ignored" in caplog.text
|
||||
|
||||
|
||||
def test_detect_languages_raises_TypeError_for_invalid_languages():
|
||||
with pytest.raises(TypeError):
|
||||
text = "This is a short sentence."
|
||||
detect_languages(text, languages="eng") == ["eng"] # type: ignore
|
||||
|
||||
|
||||
def test_apply_lang_metadata_has_no_warning_for_PageBreak(caplog: LogCaptureFixture):
|
||||
elements = [NarrativeText("Sample text."), PageBreak("")]
|
||||
elements = list(
|
||||
apply_lang_metadata(
|
||||
elements=elements,
|
||||
languages=["auto"],
|
||||
detect_language_per_element=True,
|
||||
),
|
||||
)
|
||||
assert "No features in text." not in [rec.message for rec in caplog.records]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("lang_in", "expected_lang"),
|
||||
[
|
||||
("en", "eng"),
|
||||
("fr", "fra"),
|
||||
],
|
||||
)
|
||||
def test_convert_language_code_to_pytesseract_lang_code(lang_in: str, expected_lang: str):
|
||||
assert expected_lang == _convert_language_code_to_pytesseract_lang_code(lang_in)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("input_ocr_langs", "expected"),
|
||||
[
|
||||
(["eng"], "eng"), # list
|
||||
('"deu"', "deu"), # extra quotation marks
|
||||
("[deu]", "deu"), # brackets
|
||||
("['deu']", "deu"), # brackets and quotation marks
|
||||
(["[deu]"], "deu"), # list, brackets and quotation marks
|
||||
(['"deu"'], "deu"), # list and quotation marks
|
||||
("deu+spa", "deu+spa"), # correct input
|
||||
],
|
||||
)
|
||||
def test_clean_ocr_languages_arg(input_ocr_langs: str, expected: str):
|
||||
assert _clean_ocr_languages_arg(input_ocr_langs) == expected
|
||||
|
||||
|
||||
def test_detect_languages_handles_spelled_out_languages():
|
||||
languages = detect_languages(text="Sample text longer than 5 words.", languages=["Spanish"])
|
||||
assert languages == ["spa"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("languages", "ocr_languages", "expected_langs"),
|
||||
[
|
||||
(["spa"], "deu", ["spa"]),
|
||||
(["spanish"], "english", ["spa"]),
|
||||
(["spa"], "[deu]", ["spa"]),
|
||||
(["spa"], '"deu"', ["spa"]),
|
||||
(["spa"], ["deu"], ["spa"]),
|
||||
(["spa"], ["[deu]"], ["spa"]),
|
||||
(["spa+deu"], "eng+deu", ["spa", "deu"]),
|
||||
],
|
||||
)
|
||||
def test_check_language_args_uses_languages_when_ocr_languages_and_languages_are_both_defined(
|
||||
languages: list[str],
|
||||
ocr_languages: list[str] | str,
|
||||
expected_langs: list[str],
|
||||
caplog: LogCaptureFixture,
|
||||
):
|
||||
returned_langs = check_language_args(
|
||||
languages=languages,
|
||||
ocr_languages=ocr_languages,
|
||||
)
|
||||
for lang in returned_langs: # type: ignore
|
||||
assert lang in expected_langs
|
||||
assert "ocr_languages" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("languages", "ocr_languages", "expected_langs"),
|
||||
[
|
||||
# raise warning and use `ocr_languages` when `languages` is empty or None
|
||||
([], "deu", ["deu"]),
|
||||
([""], '"deu"', ["deu"]),
|
||||
([""], "deu", ["deu"]),
|
||||
([""], "[deu]", ["deu"]),
|
||||
],
|
||||
)
|
||||
def test_check_language_args_uses_ocr_languages_when_languages_is_empty_or_None(
|
||||
languages: list[str],
|
||||
ocr_languages: str,
|
||||
expected_langs: list[str],
|
||||
caplog: LogCaptureFixture,
|
||||
):
|
||||
returned_langs = check_language_args(languages=languages, ocr_languages=ocr_languages)
|
||||
for lang in returned_langs: # type: ignore
|
||||
assert lang in expected_langs
|
||||
assert "ocr_languages" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("languages", "ocr_languages"),
|
||||
[
|
||||
([], None), # how check_language_args is called from auto.partition()
|
||||
([""], None),
|
||||
],
|
||||
)
|
||||
def test_check_language_args_returns_None(
|
||||
languages: list[str],
|
||||
ocr_languages: None,
|
||||
):
|
||||
returned_langs = check_language_args(languages=languages, ocr_languages=ocr_languages)
|
||||
assert returned_langs is None
|
||||
|
||||
|
||||
def test_check_language_args_returns_auto():
|
||||
assert check_language_args(languages=["eng", "spa", "auto"], ocr_languages=None) == ["auto"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("languages", "ocr_languages"),
|
||||
[
|
||||
([], ["auto"]),
|
||||
([""], "eng+auto"),
|
||||
],
|
||||
)
|
||||
def test_check_language_args_raises_error_when_ocr_languages_contains_auto(
|
||||
languages: list[str],
|
||||
ocr_languages: str | list[str],
|
||||
):
|
||||
with pytest.raises(ValueError):
|
||||
check_language_args(
|
||||
languages=languages,
|
||||
ocr_languages=ocr_languages,
|
||||
)
|
||||
@@ -0,0 +1,484 @@
|
||||
"""Test-suite for `unstructured.partition.common.metadata` module."""
|
||||
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import datetime as dt
|
||||
import os
|
||||
import pathlib
|
||||
from typing import Any, Callable
|
||||
|
||||
import pytest
|
||||
|
||||
from unstructured.documents.elements import (
|
||||
CheckBox,
|
||||
Element,
|
||||
ElementMetadata,
|
||||
FigureCaption,
|
||||
Header,
|
||||
ListItem,
|
||||
NarrativeText,
|
||||
Text,
|
||||
Title,
|
||||
)
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.common.metadata import (
|
||||
_assign_hash_ids,
|
||||
apply_metadata,
|
||||
get_last_modified_date,
|
||||
set_element_hierarchy,
|
||||
)
|
||||
|
||||
# ================================================================================================
|
||||
# LAST-MODIFIED
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class Describe_get_last_modified_date:
|
||||
def it_gets_the_modified_time_of_a_file_identified_by_a_path(self, tmp_path: pathlib.Path):
|
||||
modified_timestamp = dt.datetime(
|
||||
year=2024, month=3, day=5, hour=17, minute=43, second=40
|
||||
).timestamp()
|
||||
file_path = tmp_path / "some_file.txt"
|
||||
file_path.write_text("abcdefg")
|
||||
os.utime(file_path, (modified_timestamp, modified_timestamp))
|
||||
|
||||
last_modified_date = get_last_modified_date(str(file_path))
|
||||
|
||||
assert last_modified_date == "2024-03-05T17:43:40"
|
||||
|
||||
def but_it_returns_None_when_there_is_no_file_at_that_path(self, tmp_path: pathlib.Path):
|
||||
file_path = tmp_path / "some_file_that_does_not_exist.txt"
|
||||
|
||||
last_modified_date = get_last_modified_date(str(file_path))
|
||||
|
||||
assert last_modified_date is None
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# ELEMENT HIERARCHY
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class Describe_set_element_hierarchy:
|
||||
def it_applies_default_ruleset(self):
|
||||
elements = [
|
||||
Title(element_id="0", text="Title0"),
|
||||
Text(element_id="1", text="Text0"),
|
||||
Header(element_id="2", text="Header0"),
|
||||
Text(element_id="3", text="Text1"),
|
||||
Title(element_id="4", text="Title1"),
|
||||
Text(element_id="5", text="Text2"),
|
||||
]
|
||||
|
||||
result = set_element_hierarchy(elements)
|
||||
|
||||
assert result[0].metadata.parent_id is None
|
||||
assert result[1].metadata.parent_id == "0" # Text0 is under Title0
|
||||
assert result[2].metadata.parent_id is None # Header0 is higher than Title0
|
||||
assert result[3].metadata.parent_id == "2" # Text1 is under Header0
|
||||
assert result[4].metadata.parent_id == "2" # Title1 is under Header0
|
||||
assert result[5].metadata.parent_id == "4" # Text2 is under Title1, which is under Header0
|
||||
|
||||
def it_applies_category_depth_when_element_category_is_the_same(self):
|
||||
elements = [
|
||||
Title(element_id="0", text="Title0", metadata=ElementMetadata(category_depth=1)),
|
||||
ListItem(element_id="1", text="ListItem0", metadata=ElementMetadata(category_depth=0)),
|
||||
ListItem(element_id="2", text="ListItem1", metadata=ElementMetadata(category_depth=1)),
|
||||
ListItem(element_id="3", text="ListItem2", metadata=ElementMetadata(category_depth=0)),
|
||||
]
|
||||
|
||||
result = set_element_hierarchy(elements)
|
||||
|
||||
assert result[0].metadata.parent_id is None
|
||||
assert result[1].metadata.parent_id == "0" # category_depth=0
|
||||
assert result[2].metadata.parent_id == "1" # category_depth=1, so it is under ListItem0
|
||||
assert result[3].metadata.parent_id == "0" # category_depth=0
|
||||
|
||||
def but_it_ignores_category_depth_when_elements_are_of_different_categories(self):
|
||||
elements = [
|
||||
Title(element_id="0", text="Title", metadata=ElementMetadata(category_depth=2)),
|
||||
Text(element_id="1", text="Text", metadata=ElementMetadata(category_depth=0)),
|
||||
Header(element_id="2", text="Header", metadata=ElementMetadata(category_depth=2)),
|
||||
Text(element_id="3", text="Text", metadata=ElementMetadata(category_depth=0)),
|
||||
ListItem(element_id="4", text="ListItem", metadata=ElementMetadata(category_depth=1)),
|
||||
NarrativeText(element_id="5", text="", metadata=ElementMetadata(category_depth=0)),
|
||||
]
|
||||
|
||||
result = set_element_hierarchy(elements)
|
||||
|
||||
assert result[0].metadata.parent_id is None
|
||||
assert result[1].metadata.parent_id == "0" # Text is under Title despite category_depth=0
|
||||
assert result[2].metadata.parent_id is None
|
||||
assert result[3].metadata.parent_id == "2" # These are under Header despite category_depth
|
||||
assert result[4].metadata.parent_id == "2"
|
||||
assert result[5].metadata.parent_id == "2"
|
||||
|
||||
def it_skips_elements_with_pre_existing_parent_id(self):
|
||||
elements = [
|
||||
Title(element_id="0", text="Title", metadata=ElementMetadata(parent_id="10")),
|
||||
Title(element_id="1", text="Title"),
|
||||
Text(element_id="2", text="Text"),
|
||||
]
|
||||
|
||||
result = set_element_hierarchy(elements)
|
||||
|
||||
# Parent ID should not change and element is skipped in figuring out other elements' parents
|
||||
assert result[0].metadata.parent_id == "10"
|
||||
assert result[1].metadata.parent_id is None
|
||||
assert result[2].metadata.parent_id == "1"
|
||||
|
||||
def it_sets_parent_id_for_each_element_in_elements(self):
|
||||
elements_to_set = [
|
||||
Title(text="Title"), # 0
|
||||
NarrativeText(text="NarrativeText"), # 1
|
||||
FigureCaption(text="FigureCaption"), # 2
|
||||
ListItem(text="ListItem"), # 3
|
||||
ListItem(text="ListItem", metadata=ElementMetadata(category_depth=1)), # 4
|
||||
ListItem(text="ListItem", metadata=ElementMetadata(category_depth=1)), # 5
|
||||
ListItem(text="ListItem"), # 6
|
||||
CheckBox(element_id="some-id-1", checked=True), # 7
|
||||
Title(text="Title 2"), # 8
|
||||
ListItem(text="ListItem"), # 9
|
||||
ListItem(text="ListItem"), # 10
|
||||
Text(text="Text"), # 11
|
||||
]
|
||||
elements = set_element_hierarchy(elements_to_set)
|
||||
|
||||
assert (
|
||||
elements[1].metadata.parent_id == elements[0].id
|
||||
), "NarrativeText should be child of Title"
|
||||
assert (
|
||||
elements[2].metadata.parent_id == elements[0].id
|
||||
), "FigureCaption should be child of Title"
|
||||
assert elements[3].metadata.parent_id == elements[0].id, "ListItem should be child of Title"
|
||||
assert elements[4].metadata.parent_id == elements[3].id, "ListItem should be child of Title"
|
||||
assert elements[5].metadata.parent_id == elements[3].id, "ListItem should be child of Title"
|
||||
assert elements[6].metadata.parent_id == elements[0].id, "ListItem should be child of Title"
|
||||
# NOTE(Hubert): moving the category field to Element, caused this to fail.
|
||||
# Checkboxes will soon be deprecated, then we can remove the test.
|
||||
# assert (
|
||||
# elements[7].metadata.parent_id is None
|
||||
# ), "CheckBox should be None, as it's not a Text based element"
|
||||
assert elements[8].metadata.parent_id is None, "Title 2 should be child of None"
|
||||
assert (
|
||||
elements[9].metadata.parent_id == elements[8].id
|
||||
), "ListItem should be child of Title 2"
|
||||
assert (
|
||||
elements[10].metadata.parent_id == elements[8].id
|
||||
), "ListItem should be child of Title 2"
|
||||
assert elements[11].metadata.parent_id == elements[8].id, "Text should be child of Title 2"
|
||||
|
||||
def it_applies_custom_rule_set(self):
|
||||
elements_to_set = [
|
||||
Header(text="Header"), # 0
|
||||
Title(text="Title"), # 1
|
||||
NarrativeText(text="NarrativeText"), # 2
|
||||
Text(text="Text"), # 3
|
||||
Title(text="Title 2"), # 4
|
||||
FigureCaption(text="FigureCaption"), # 5
|
||||
]
|
||||
|
||||
custom_rule_set = {
|
||||
"Header": ["Title", "Text"],
|
||||
"Title": ["NarrativeText", "UncategorizedText", "FigureCaption"],
|
||||
}
|
||||
|
||||
elements = set_element_hierarchy(
|
||||
elements=elements_to_set,
|
||||
ruleset=custom_rule_set,
|
||||
)
|
||||
|
||||
assert elements[1].metadata.parent_id == elements[0].id, "Title should be child of Header"
|
||||
assert (
|
||||
elements[2].metadata.parent_id == elements[1].id
|
||||
), "NarrativeText should be child of Title"
|
||||
assert elements[3].metadata.parent_id == elements[1].id, "Text should be child of Title"
|
||||
assert elements[4].metadata.parent_id == elements[0].id, "Title 2 should be child of Header"
|
||||
assert (
|
||||
elements[5].metadata.parent_id == elements[4].id
|
||||
), "FigureCaption should be child of Title 2"
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# APPLY METADATA DECORATOR
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class Describe_apply_metadata:
|
||||
"""Unit-test suite for `unstructured.partition.common.metadata.apply_metadata()` decorator."""
|
||||
|
||||
# -- unique-ify elements and metadata ---------------------------------
|
||||
|
||||
def it_produces_unique_elements_and_metadata_when_input_reuses_element_instances(self):
|
||||
element = Text(text="Element", metadata=ElementMetadata(filename="foo.bar", page_number=1))
|
||||
|
||||
def fake_partitioner(**kwargs: Any) -> list[Element]:
|
||||
return [element, element, element]
|
||||
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition()
|
||||
|
||||
# -- all elements are unique instances --
|
||||
assert len({id(e) for e in elements}) == len(elements)
|
||||
# -- all metadatas are unique instances --
|
||||
assert len({id(e.metadata) for e in elements}) == len(elements)
|
||||
|
||||
def and_it_produces_unique_elements_and_metadata_when_input_reuses_metadata_instances(self):
|
||||
metadata = ElementMetadata(filename="foo.bar", page_number=1)
|
||||
|
||||
def fake_partitioner(**kwargs: Any) -> list[Element]:
|
||||
return [
|
||||
Text(text="foo", metadata=metadata),
|
||||
Text(text="bar", metadata=metadata),
|
||||
Text(text="baz", metadata=metadata),
|
||||
]
|
||||
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition()
|
||||
|
||||
# -- all elements are unique instances --
|
||||
assert len({id(e) for e in elements}) == len(elements)
|
||||
# -- all metadatas are unique instances --
|
||||
assert len({id(e.metadata) for e in elements}) == len(elements)
|
||||
|
||||
# -- unique-ids -------------------------------------------------------
|
||||
|
||||
def it_assigns_hash_element_ids_when_unique_ids_arg_is_not_specified(
|
||||
self, fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition()
|
||||
elements_2 = partition()
|
||||
|
||||
# -- SHA1 hash is 32 characters long, no hyphens --
|
||||
assert all(len(e.id) == 32 for e in elements)
|
||||
assert all("-" not in e.id for e in elements)
|
||||
# -- SHA1 hashes are deterministic --
|
||||
assert all(e.id == e2.id for e, e2 in zip(elements, elements_2))
|
||||
|
||||
def it_assigns_hash_element_ids_when_unique_ids_arg_is_False(
|
||||
self, fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition(unique_element_ids=False)
|
||||
elements_2 = partition(unique_element_ids=False)
|
||||
|
||||
# -- SHA1 hash is 32 characters long, no hyphens --
|
||||
assert all(len(e.id) == 32 for e in elements)
|
||||
assert all("-" not in e.id for e in elements)
|
||||
# -- SHA1 hashes are deterministic --
|
||||
assert all(e.id == e2.id for e, e2 in zip(elements, elements_2))
|
||||
|
||||
def it_leaves_UUID_element_ids_when_unique_ids_arg_is_True(
|
||||
self, fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition(unique_element_ids=True)
|
||||
elements_2 = partition(unique_element_ids=True)
|
||||
|
||||
# -- UUID is 36 characters long with four hyphens --
|
||||
assert all(len(e.id) == 36 for e in elements)
|
||||
assert all(e.id.count("-") == 4 for e in elements)
|
||||
# -- UUIDs are non-deterministic, different every time --
|
||||
assert all(e.id != e2.id for e, e2 in zip(elements, elements_2))
|
||||
|
||||
# -- parent-id --------------------------------------------------------
|
||||
|
||||
def it_computes_and_assigns_parent_id(self, fake_partitioner: Callable[..., list[Element]]):
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition()
|
||||
|
||||
title = elements[0]
|
||||
assert title.metadata.category_depth == 1
|
||||
narr_text = elements[1]
|
||||
assert narr_text.metadata.parent_id == title.id
|
||||
|
||||
# -- languages --------------------------------------------------------
|
||||
|
||||
def it_applies_language_metadata(self, fake_partitioner: Callable[..., list[Element]]):
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition(languages=["auto"], detect_language_per_element=True)
|
||||
|
||||
assert all(e.metadata.languages == ["eng"] for e in elements)
|
||||
|
||||
# -- filetype (MIME-type) ---------------------------------------------
|
||||
|
||||
def it_assigns_the_value_of_a_metadata_file_type_arg_when_there_is_one(
|
||||
self, fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
"""A `metadata_file_type` arg overrides the file-type specified in the decorator.
|
||||
|
||||
This is used for example by a delegating partitioner to preserve the original file-type in
|
||||
the metadata, like EPUB instead of the HTML that partitioner converts the .epub file to.
|
||||
"""
|
||||
partition = apply_metadata(file_type=FileType.DOCX)(fake_partitioner)
|
||||
|
||||
elements = partition(metadata_file_type=FileType.ODT)
|
||||
|
||||
assert all(
|
||||
e.metadata.filetype == "application/vnd.oasis.opendocument.text" for e in elements
|
||||
)
|
||||
|
||||
def and_it_assigns_the_decorator_file_type_when_the_metadata_file_type_arg_is_omitted(
|
||||
self, fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
"""The `file_type=...` decorator arg is the "normal" way to specify the file-type.
|
||||
|
||||
This is used for principal (non-delegating) partitioners.
|
||||
"""
|
||||
partition = apply_metadata(file_type=FileType.DOCX)(fake_partitioner)
|
||||
|
||||
elements = partition()
|
||||
|
||||
DOCX_MIME_TYPE = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||||
assert all(e.metadata.filetype == DOCX_MIME_TYPE for e in elements)
|
||||
|
||||
def and_it_does_not_assign_file_type_metadata_when_both_are_omitted(
|
||||
self, fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
"""A partitioner can elect to assign `.metadata.filetype` for itself.
|
||||
|
||||
This is done in `partition_image()` for example where the same partitioner is used for
|
||||
multiple file-types.
|
||||
"""
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition()
|
||||
|
||||
assert all(e.metadata.filetype == "image/jpeg" for e in elements)
|
||||
|
||||
# -- filename ---------------------------------------------------------
|
||||
|
||||
def it_uses_metadata_filename_arg_value_when_present(
|
||||
self, fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
"""A `metadata_filename` arg overrides all other sources."""
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition(metadata_filename="a/b/c.xyz")
|
||||
|
||||
assert all(e.metadata.filename == "c.xyz" for e in elements)
|
||||
assert all(e.metadata.file_directory == "a/b" for e in elements)
|
||||
|
||||
def and_it_uses_filename_arg_value_when_metadata_filename_arg_not_present(
|
||||
self, fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition(filename="a/b/c.xyz")
|
||||
|
||||
assert all(e.metadata.filename == "c.xyz" for e in elements)
|
||||
assert all(e.metadata.file_directory == "a/b" for e in elements)
|
||||
|
||||
def and_it_does_not_assign_filename_metadata_when_neither_are_present(
|
||||
self, fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition()
|
||||
|
||||
assert all(e.metadata.filename == "image.jpeg" for e in elements)
|
||||
assert all(e.metadata.file_directory == "x/y/images" for e in elements)
|
||||
|
||||
# -- last_modified ----------------------------------------------------
|
||||
|
||||
def it_uses_metadata_last_modified_arg_value_when_present(
|
||||
self, fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
"""A `metadata_last_modified` arg overrides all other sources."""
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
metadata_last_modified = "2024-09-26T15:17:53"
|
||||
|
||||
elements = partition(metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
@pytest.mark.parametrize("kwargs", [{}, {"metadata_last_modified": None}])
|
||||
def but_it_does_not_update_last_modified_when_metadata_last_modified_arg_absent_or_None(
|
||||
self, kwargs: dict[str, Any], fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition(**kwargs)
|
||||
|
||||
assert all(e.metadata.last_modified == "2020-01-06T05:07:03" for e in elements)
|
||||
|
||||
# -- url --------------------------------------------------------------
|
||||
|
||||
def it_assigns_url_metadata_field_when_url_arg_is_present(
|
||||
self, fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition(url="https://adobe.com/stock/54321")
|
||||
|
||||
assert all(e.metadata.url == "https://adobe.com/stock/54321" for e in elements)
|
||||
|
||||
def and_it_does_not_assign_url_metadata_when_url_arg_is_not_present(
|
||||
self, fake_partitioner: Callable[..., list[Element]]
|
||||
):
|
||||
partition = apply_metadata()(fake_partitioner)
|
||||
|
||||
elements = partition()
|
||||
|
||||
assert all(e.metadata.url == "http://images.com" for e in elements)
|
||||
|
||||
# -- fixtures --------------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture
|
||||
def fake_partitioner(self) -> Callable[..., list[Element]]:
|
||||
def fake_partitioner(**kwargs: Any) -> list[Element]:
|
||||
title = Title("Introduction")
|
||||
title.metadata.category_depth = 1
|
||||
title.metadata.file_directory = "x/y/images"
|
||||
title.metadata.filename = "image.jpeg"
|
||||
title.metadata.filetype = "image/jpeg"
|
||||
title.metadata.last_modified = "2020-01-06T05:07:03"
|
||||
title.metadata.url = "http://images.com"
|
||||
|
||||
narr_text = NarrativeText("To understand bar you must first understand foo.")
|
||||
narr_text.metadata.file_directory = "x/y/images"
|
||||
narr_text.metadata.filename = "image.jpeg"
|
||||
narr_text.metadata.filetype = "image/jpeg"
|
||||
narr_text.metadata.last_modified = "2020-01-06T05:07:03"
|
||||
narr_text.metadata.url = "http://images.com"
|
||||
|
||||
return [title, narr_text]
|
||||
|
||||
return fake_partitioner
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# HASH IDS
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
def test_assign_hash_ids_produces_unique_and_deterministic_SHA1_ids_even_for_duplicate_elements():
|
||||
elements: list[Element] = [
|
||||
Text(text="Element", metadata=ElementMetadata(filename="foo.bar", page_number=1)),
|
||||
Text(text="Element", metadata=ElementMetadata(filename="foo.bar", page_number=1)),
|
||||
Text(text="Element", metadata=ElementMetadata(filename="foo.bar", page_number=1)),
|
||||
]
|
||||
# -- default ids are UUIDs --
|
||||
assert all(len(e.id) == 36 for e in elements)
|
||||
|
||||
elements = _assign_hash_ids(copy.deepcopy(elements))
|
||||
elements_2 = _assign_hash_ids(copy.deepcopy(elements))
|
||||
|
||||
ids = [e.id for e in elements]
|
||||
# -- ids are now SHA1 --
|
||||
assert all(len(e.id) == 32 for e in elements)
|
||||
# -- each id is unique --
|
||||
assert len(ids) == len(set(ids))
|
||||
# -- ids are deterministic, same value is computed each time --
|
||||
assert all(e.id == e2.id for e, e2 in zip(elements, elements_2))
|
||||
@@ -0,0 +1,18 @@
|
||||
import pytest
|
||||
|
||||
from unstructured.partition.utils.constants import OCR_AGENT_PADDLE, OCR_AGENT_TESSERACT
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_ocr_get_instance(mocker):
|
||||
"""Fixture that mocks OCRAgent.get_instance to prevent real OCR agent instantiation."""
|
||||
|
||||
def mock_get_instance(ocr_agent_module, language):
|
||||
if ocr_agent_module in (OCR_AGENT_TESSERACT, OCR_AGENT_PADDLE):
|
||||
return mocker.MagicMock()
|
||||
else:
|
||||
raise ValueError(f"Unknown OCR agent: {ocr_agent_module}")
|
||||
|
||||
from unstructured.partition.pdf_image.ocr import OCRAgent
|
||||
|
||||
return mocker.patch.object(OCRAgent, "get_instance", side_effect=mock_get_instance)
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,489 @@
|
||||
# pyright: reportPrivateUsage=false, reportUnknownMemberType=false, reportOptionalMemberAccess=false
|
||||
# pyright: reportAttributeAccessIssue=false, reportUnknownLambdaType=false
|
||||
from collections import defaultdict
|
||||
from typing import Any, Optional
|
||||
|
||||
import pytest
|
||||
from bs4 import BeautifulSoup, Tag
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from unstructured.documents.elements import Element, ElementMetadata, ElementType
|
||||
from unstructured.partition.html.convert import (
|
||||
HTML_PARSER,
|
||||
ElementHtml,
|
||||
ImageElementHtml,
|
||||
LinkElementHtml,
|
||||
ListItemElementHtml,
|
||||
TableElementHtml,
|
||||
TextElementHtml,
|
||||
TitleElementHtml,
|
||||
UnorderedListElementHtml,
|
||||
_elements_to_html_tags,
|
||||
_elements_to_html_tags_by_page,
|
||||
_elements_to_html_tags_by_parent,
|
||||
_group_element_children,
|
||||
elements_to_html,
|
||||
group_elements_by_page,
|
||||
)
|
||||
|
||||
|
||||
class MockElement(Element):
|
||||
def __init__(
|
||||
self,
|
||||
text: str = "",
|
||||
metadata: Optional[ElementMetadata] = None,
|
||||
category: str = "",
|
||||
id: str = "",
|
||||
) -> None:
|
||||
self.text = text
|
||||
self.metadata = metadata or ElementMetadata()
|
||||
self.category = category
|
||||
self._element_id = id
|
||||
|
||||
|
||||
class MockElementMetadata(ElementMetadata):
|
||||
def __init__(
|
||||
self,
|
||||
text_as_html: Optional[str] = None,
|
||||
category_depth: Optional[int] = None,
|
||||
image_base64: Optional[str] = None,
|
||||
image_mime_type: Optional[str] = None,
|
||||
url: Optional[str] = None,
|
||||
parent_id: Optional[str] = None,
|
||||
page_number: Optional[int] = None,
|
||||
) -> None:
|
||||
self.text_as_html = text_as_html
|
||||
self.category_depth = category_depth
|
||||
self.image_base64 = image_base64
|
||||
self.image_mime_type = image_mime_type
|
||||
self.url = url
|
||||
self.parent_id = parent_id
|
||||
self.page_number = page_number
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_element() -> MockElement:
|
||||
metadata = MockElementMetadata(text_as_html="<p>Test Text</p>")
|
||||
return MockElement(text="Test Text", metadata=metadata, category="test-category", id="test-id")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def element_html(mock_element: MockElement) -> ElementHtml:
|
||||
return ElementHtml(mock_element)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def title_element_html(mock_element: MockElement) -> TitleElementHtml:
|
||||
metadata = MockElementMetadata(text_as_html="<p>Test HTML</p>")
|
||||
MockElement(text="Test Text", metadata=metadata, category="test-category", id="test-id")
|
||||
return TitleElementHtml(mock_element)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def image_element_html(mock_element: MockElement) -> ImageElementHtml:
|
||||
return ImageElementHtml(mock_element)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def table_element_html(mock_element: MockElement) -> TableElementHtml:
|
||||
return TableElementHtml(mock_element)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def link_element_html(mock_element: MockElement) -> LinkElementHtml:
|
||||
return LinkElementHtml(mock_element)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def unordered_list_element_html(mock_element: MockElement) -> UnorderedListElementHtml:
|
||||
return UnorderedListElementHtml(mock_element)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def elements_html() -> list[ElementHtml]:
|
||||
return [
|
||||
ListItemElementHtml(
|
||||
MockElement(
|
||||
text="Test List Item",
|
||||
metadata=MockElementMetadata(page_number=1),
|
||||
category=ElementType.LIST_ITEM,
|
||||
id="test-element-1",
|
||||
)
|
||||
),
|
||||
TextElementHtml(
|
||||
MockElement(
|
||||
text="Test Text",
|
||||
metadata=MockElementMetadata(page_number=None),
|
||||
category=ElementType.TEXT,
|
||||
id="test-element-2",
|
||||
)
|
||||
),
|
||||
TextElementHtml(
|
||||
MockElement(
|
||||
text="Test Text",
|
||||
metadata=MockElementMetadata(page_number=2),
|
||||
category=ElementType.TEXT,
|
||||
id="test-element-3",
|
||||
)
|
||||
),
|
||||
ListItemElementHtml(
|
||||
MockElement(
|
||||
text="Test List Item",
|
||||
metadata=MockElementMetadata(parent_id="test-element-3", page_number=2),
|
||||
category=ElementType.LIST_ITEM,
|
||||
id="test-element-4",
|
||||
)
|
||||
),
|
||||
ListItemElementHtml(
|
||||
MockElement(
|
||||
text="Test List Item",
|
||||
metadata=MockElementMetadata(parent_id="test-element-3", page_number=2),
|
||||
category=ElementType.LIST_ITEM,
|
||||
id="test-element-5",
|
||||
)
|
||||
),
|
||||
TextElementHtml(
|
||||
MockElement(
|
||||
text="Test Text",
|
||||
metadata=MockElementMetadata(page_number=3),
|
||||
category=ElementType.TEXT,
|
||||
id="test-element-6",
|
||||
)
|
||||
),
|
||||
ListItemElementHtml(
|
||||
MockElement(
|
||||
text="Test List Item Other",
|
||||
metadata=MockElementMetadata(parent_id="test-element-6", page_number=3),
|
||||
category=ElementType.LIST_ITEM_OTHER,
|
||||
id="test-element-7",
|
||||
)
|
||||
),
|
||||
ListItemElementHtml(
|
||||
MockElement(
|
||||
text="Test List Item",
|
||||
metadata=MockElementMetadata(parent_id="test-element-7", page_number=3),
|
||||
category=ElementType.LIST_ITEM,
|
||||
id="test-element-8",
|
||||
)
|
||||
),
|
||||
ListItemElementHtml(
|
||||
MockElement(
|
||||
text="Test List Item Other",
|
||||
metadata=MockElementMetadata(parent_id="test-element-6", page_number=3),
|
||||
category=ElementType.LIST_ITEM_OTHER,
|
||||
id="test-element-9",
|
||||
)
|
||||
),
|
||||
TextElementHtml(
|
||||
MockElement(
|
||||
text="Test Text",
|
||||
metadata=MockElementMetadata(parent_id="test-element-6", page_number=3),
|
||||
category=ElementType.TEXT,
|
||||
id="test-element-10",
|
||||
)
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def elements(elements_html: list[ElementHtml]) -> list[Element]:
|
||||
return [el.element for el in elements_html]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def elements_small() -> list[Element]:
|
||||
return [
|
||||
MockElement(
|
||||
text="Test Text 1",
|
||||
category=ElementType.TEXT,
|
||||
id="test-element-1",
|
||||
metadata=MockElementMetadata(page_number=1),
|
||||
),
|
||||
MockElement(
|
||||
text="Test Text 2",
|
||||
category=ElementType.TEXT,
|
||||
id="test-element-2",
|
||||
metadata=MockElementMetadata(page_number=2),
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_inject_html_element_content(element_html: ElementHtml) -> None:
|
||||
soup = BeautifulSoup("", HTML_PARSER)
|
||||
tag = soup.new_tag("div")
|
||||
element_html._inject_html_element_content(tag)
|
||||
assert tag.string == "Test Text"
|
||||
|
||||
|
||||
def test_get_text_as_html(element_html: ElementHtml) -> None:
|
||||
tag = element_html.get_text_as_html()
|
||||
assert isinstance(tag, Tag)
|
||||
assert tag.name == "p"
|
||||
assert tag.string == "Test Text"
|
||||
|
||||
|
||||
def test_get_children_html(element_html: ElementHtml) -> None:
|
||||
soup = BeautifulSoup("", HTML_PARSER)
|
||||
parent_tag = soup.new_tag("div")
|
||||
child_element = MockElement(text="Child Text")
|
||||
child_element_html = ElementHtml(child_element)
|
||||
element_html.set_children([child_element_html])
|
||||
result_tag = element_html._get_children_html(soup, parent_tag)
|
||||
assert result_tag.name == "div"
|
||||
assert len(result_tag.contents) == 2
|
||||
assert result_tag.contents[1].string == "Child Text"
|
||||
|
||||
|
||||
def test_get_html_element_text_as_html(element_html: ElementHtml) -> None:
|
||||
tag = element_html.get_html_element()
|
||||
assert isinstance(tag, Tag)
|
||||
assert tag.name == "p"
|
||||
assert tag.string == "Test Text"
|
||||
assert tag["class"] == "test-category"
|
||||
assert tag["id"] == "test-id"
|
||||
|
||||
|
||||
def test_get_html_element_no_text_as_html(element_html: ElementHtml) -> None:
|
||||
element_html.element.metadata.text_as_html = None
|
||||
tag = element_html.get_html_element()
|
||||
assert isinstance(tag, Tag)
|
||||
assert tag.name == "div"
|
||||
assert tag.string == "Test Text"
|
||||
assert tag["class"] == "test-category"
|
||||
assert tag["id"] == "test-id"
|
||||
|
||||
|
||||
def test_set_children(element_html: ElementHtml) -> None:
|
||||
child_element = MockElement(text="Child Text")
|
||||
child_element_html = ElementHtml(child_element)
|
||||
element_html.set_children([child_element_html])
|
||||
assert len(element_html.children) == 1
|
||||
assert element_html.children[0].element.text == "Child Text"
|
||||
|
||||
|
||||
def test_title_element_html_tag(title_element_html: TitleElementHtml) -> None:
|
||||
assert title_element_html.html_tag == "h1"
|
||||
|
||||
|
||||
def test_image_element_html_content(image_element_html: ImageElementHtml) -> None:
|
||||
soup = BeautifulSoup("", HTML_PARSER)
|
||||
tag = soup.new_tag("img")
|
||||
image_element_html._inject_html_element_content(tag)
|
||||
assert tag["alt"] == "Test Text"
|
||||
|
||||
|
||||
def test_image_element_html_content_with_base64(image_element_html: ImageElementHtml) -> None:
|
||||
image_element_html.element.metadata.image_base64 = "base64data"
|
||||
image_element_html.element.metadata.image_mime_type = "image/png"
|
||||
soup = BeautifulSoup("", HTML_PARSER)
|
||||
tag = soup.new_tag("img")
|
||||
image_element_html._inject_html_element_content(tag)
|
||||
assert tag["src"] == "data:image/png;base64,base64data"
|
||||
assert tag["alt"] == "Test Text"
|
||||
|
||||
|
||||
def test_table_element_html_attrs(table_element_html: TableElementHtml) -> None:
|
||||
soup = BeautifulSoup("", HTML_PARSER)
|
||||
tag = soup.new_tag("table")
|
||||
table_element_html._inject_html_element_attrs(tag)
|
||||
assert tag["style"] == "border: 1px solid black; border-collapse: collapse;"
|
||||
|
||||
|
||||
def test_link_element_html_attrs(link_element_html: LinkElementHtml) -> None:
|
||||
link_element_html.element.metadata.url = "http://example.com"
|
||||
soup = BeautifulSoup("", HTML_PARSER)
|
||||
tag = soup.new_tag("a")
|
||||
link_element_html._inject_html_element_attrs(tag)
|
||||
assert tag["href"] == "http://example.com"
|
||||
|
||||
|
||||
def test_unordered_list_element_html(unordered_list_element_html: UnorderedListElementHtml) -> None:
|
||||
soup = BeautifulSoup("", HTML_PARSER)
|
||||
tag = soup.new_tag("ul")
|
||||
child_element = MockElement(text="Child Text")
|
||||
child_element_html = ListItemElementHtml(child_element)
|
||||
unordered_list_element_html.set_children([child_element_html])
|
||||
result_tag = unordered_list_element_html._get_children_html(soup, tag)
|
||||
assert result_tag.name == "ul"
|
||||
assert len(result_tag.contents) == 1
|
||||
assert result_tag.contents[0].name == "li"
|
||||
assert result_tag.contents[0].string == "Child Text"
|
||||
|
||||
|
||||
def test_group_element_children(elements_html: list[ElementHtml]) -> None:
|
||||
grouped_children = _group_element_children(elements_html)
|
||||
assert len(grouped_children) == 7
|
||||
assert len(grouped_children[0].children) == 1
|
||||
assert grouped_children[0].children[0].element.category == ElementType.LIST_ITEM
|
||||
assert len(grouped_children[3].children) == 2
|
||||
assert grouped_children[3].children[0].element.category == ElementType.LIST_ITEM
|
||||
assert grouped_children[3].children[1].element.category == ElementType.LIST_ITEM
|
||||
assert len(grouped_children[5].children) == 3
|
||||
assert grouped_children[5].children[0].element.category == ElementType.LIST_ITEM_OTHER
|
||||
assert grouped_children[5].children[1].element.category == ElementType.LIST_ITEM
|
||||
assert grouped_children[5].children[2].element.category == ElementType.LIST_ITEM_OTHER
|
||||
|
||||
|
||||
def test_elements_to_html_tags_by_parent(
|
||||
mocker: MockerFixture, elements_html: list[ElementHtml]
|
||||
) -> None:
|
||||
mocker.patch(
|
||||
"unstructured.partition.html.convert._group_element_children",
|
||||
side_effect=lambda children: children,
|
||||
)
|
||||
result = _elements_to_html_tags_by_parent(elements_html)
|
||||
assert len(result) == 4
|
||||
assert result[0].element.id == "test-element-1"
|
||||
assert len(result[0].children) == 0
|
||||
assert result[1].element.id == "test-element-2"
|
||||
assert len(result[1].children) == 0
|
||||
assert result[2].element.id == "test-element-3"
|
||||
assert len(result[2].children) == 2
|
||||
assert result[2].children[0].element.id == "test-element-4"
|
||||
assert result[2].children[1].element.id == "test-element-5"
|
||||
assert result[3].element.id == "test-element-6"
|
||||
assert len(result[3].children) == 3
|
||||
assert result[3].children[0].element.id == "test-element-7"
|
||||
assert len(result[3].children[0].children) == 1
|
||||
assert result[3].children[0].children[0].element.id == "test-element-8"
|
||||
assert result[3].children[1].element.id == "test-element-9"
|
||||
assert result[3].children[2].element.id == "test-element-10"
|
||||
|
||||
|
||||
def test_elements_to_html_tags(mocker: MockerFixture, elements: list[Element]) -> None:
|
||||
def _mock_get_html_element(self: ElementHtml, **kwargs: Any):
|
||||
return BeautifulSoup(f"<div>{self.element.id}</div>", HTML_PARSER).find()
|
||||
|
||||
mocker.patch(
|
||||
"unstructured.partition.html.convert._elements_to_html_tags_by_parent",
|
||||
side_effect=lambda elements: elements,
|
||||
)
|
||||
mocker.patch(
|
||||
"unstructured.partition.html.convert.ElementHtml.get_html_element",
|
||||
side_effect=_mock_get_html_element,
|
||||
autospec=True,
|
||||
)
|
||||
result = _elements_to_html_tags(elements)
|
||||
assert len(result) == 10
|
||||
assert all(isinstance(tag, Tag) for tag in result)
|
||||
for i, el in enumerate(result, start=1):
|
||||
assert el.string == f"test-element-{i}"
|
||||
|
||||
|
||||
def test_elements_to_html_tags_by_page(mocker: MockerFixture, elements: list[Element]) -> None:
|
||||
def _mock_elements_to_html_tags(elements: list[Element], _: bool):
|
||||
return [
|
||||
BeautifulSoup(f"<div>{element.id}</div>", HTML_PARSER).find() for element in elements
|
||||
]
|
||||
|
||||
def _mock_group_elements_by_page(elements: list[Element]) -> list[list[Element]]:
|
||||
pages_dict: defaultdict[int, list[Element]] = defaultdict(list)
|
||||
for element in elements:
|
||||
if element.metadata.page_number is not None:
|
||||
pages_dict[element.metadata.page_number].append(element)
|
||||
return list(pages_dict.values())
|
||||
|
||||
mocker.patch(
|
||||
"unstructured.partition.html.convert._elements_to_html_tags",
|
||||
side_effect=_mock_elements_to_html_tags,
|
||||
)
|
||||
mocker.patch(
|
||||
"unstructured.partition.html.convert.group_elements_by_page",
|
||||
side_effect=_mock_group_elements_by_page,
|
||||
)
|
||||
result = _elements_to_html_tags_by_page(elements)
|
||||
assert len(result) == 3
|
||||
assert all(isinstance(tag, Tag) for tag in result)
|
||||
assert result[0].name == "div"
|
||||
assert result[0]["data-page_number"] == 1
|
||||
assert len(result[0].contents) == 1
|
||||
assert result[0].contents[0].string == "test-element-1"
|
||||
assert result[1].name == "div"
|
||||
assert result[1]["data-page_number"] == 2
|
||||
assert len(result[1].contents) == 3
|
||||
for i, el in enumerate(result[1].contents, start=3):
|
||||
assert el.string == f"test-element-{i}"
|
||||
assert result[2]["data-page_number"] == 3
|
||||
assert len(result[2].contents) == 5
|
||||
for i, el in enumerate(result[2].contents, start=6):
|
||||
assert el.string == f"test-element-{i}"
|
||||
|
||||
|
||||
def test_group_elements_by_page(caplog: pytest.LogCaptureFixture, elements: list[Element]) -> None:
|
||||
result = group_elements_by_page(elements)
|
||||
assert len(result) == 3
|
||||
assert len(result[0]) == 1
|
||||
assert len(result[1]) == 3
|
||||
assert len(result[2]) == 5
|
||||
assert result[0][0].id == "test-element-1"
|
||||
for i, el in enumerate(result[1], start=3):
|
||||
assert el.id == f"test-element-{i}"
|
||||
for i, el in enumerate(result[2], start=6):
|
||||
assert el.id == f"test-element-{i}"
|
||||
assert "Page number is not set for an element test-element-2. Skipping." in caplog.text
|
||||
|
||||
|
||||
def test_elements_to_html_no_group_by_page(
|
||||
mocker: MockerFixture, elements_small: list[Element]
|
||||
) -> None:
|
||||
def _mock_elements_to_html_tags(elements: list[Element], _: bool):
|
||||
return [
|
||||
BeautifulSoup(f"<div>{element.id}</div>", HTML_PARSER).find() for element in elements
|
||||
]
|
||||
|
||||
mock_elements_to_html_tags = mocker.patch(
|
||||
"unstructured.partition.html.convert._elements_to_html_tags",
|
||||
side_effect=_mock_elements_to_html_tags,
|
||||
)
|
||||
mock_elements_to_html_tags_by_page = mocker.patch(
|
||||
"unstructured.partition.html.convert._elements_to_html_tags_by_page",
|
||||
)
|
||||
|
||||
result = elements_to_html(elements_small, exclude_binary_image_data=True, no_group_by_page=True)
|
||||
assert "<div>\n test-element-1\n </div>" in result
|
||||
assert "<div>\n test-element-2\n </div>" in result
|
||||
assert "data-page_number" not in result
|
||||
mock_elements_to_html_tags.assert_called_once_with(elements_small, True)
|
||||
mock_elements_to_html_tags_by_page.assert_not_called()
|
||||
|
||||
|
||||
def test_elements_to_html_group_by_page(
|
||||
mocker: MockerFixture, elements_small: list[Element]
|
||||
) -> None:
|
||||
def _mock_elements_to_html_tags_by_page(elements: list[Element], _: bool):
|
||||
return [
|
||||
BeautifulSoup(
|
||||
f"<div data-page_number='{element.metadata.page_number}'>{element.id}</div>",
|
||||
HTML_PARSER,
|
||||
).find()
|
||||
for element in elements
|
||||
]
|
||||
|
||||
mock_elements_to_html_tags_by_page = mocker.patch(
|
||||
"unstructured.partition.html.convert._elements_to_html_tags_by_page",
|
||||
side_effect=_mock_elements_to_html_tags_by_page,
|
||||
)
|
||||
mock_elements_to_html_tags = mocker.patch(
|
||||
"unstructured.partition.html.convert._elements_to_html_tags"
|
||||
)
|
||||
|
||||
result = elements_to_html(
|
||||
elements_small, exclude_binary_image_data=True, no_group_by_page=False
|
||||
)
|
||||
soup = BeautifulSoup(result, HTML_PARSER)
|
||||
assert soup.find("div", {"data-page_number": "1"}).string.strip() == "test-element-1"
|
||||
assert soup.find("div", {"data-page_number": "2"}).string.strip() == "test-element-2"
|
||||
mock_elements_to_html_tags_by_page.assert_called_once_with(elements_small, True)
|
||||
mock_elements_to_html_tags.assert_not_called()
|
||||
|
||||
|
||||
def test_elements_to_html_invalid_html_template(
|
||||
mocker: MockerFixture, elements: list[Element]
|
||||
) -> None:
|
||||
mocker.patch(
|
||||
"unstructured.partition.html.convert.HTML_TEMPLATE",
|
||||
"<html><head><title>Test</title></head></html>",
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Body tag not found in the HTML template"):
|
||||
elements_to_html(elements)
|
||||
@@ -0,0 +1,753 @@
|
||||
from typing import Optional, Type
|
||||
|
||||
import pytest
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
from unstructured.documents.ontology import (
|
||||
Checkbox,
|
||||
Form,
|
||||
FormFieldValue,
|
||||
Image,
|
||||
OntologyElement,
|
||||
Page,
|
||||
RadioButton,
|
||||
)
|
||||
from unstructured.partition.html.html_utils import indent_html
|
||||
from unstructured.partition.html.transformations import RECURSION_LIMIT, parse_html_to_ontology
|
||||
|
||||
|
||||
def _wrap_with_body(html: str) -> str:
|
||||
return f'<body class="Document">{html}</body>'
|
||||
|
||||
|
||||
def remove_all_ids(html_str):
|
||||
soup = BeautifulSoup(html_str, "html.parser")
|
||||
for tag in soup.find_all(True):
|
||||
if tag.has_attr("id"):
|
||||
del tag["id"]
|
||||
return str(soup)
|
||||
|
||||
|
||||
def test_parsing_header_and_footer_into_correct_ontologyelement():
|
||||
input_html = """
|
||||
<div class="Page">
|
||||
<header class="Header">
|
||||
this is a header
|
||||
</header>
|
||||
<footer class="Footer">
|
||||
this is a footer
|
||||
</footer>
|
||||
</div>
|
||||
"""
|
||||
page = parse_html_to_ontology(input_html)
|
||||
assert len(page.children) == 2
|
||||
header, footer = page.children
|
||||
assert header.text == "this is a header"
|
||||
assert header.html_tag_name == "header"
|
||||
assert footer.text == "this is a footer"
|
||||
assert footer.html_tag_name == "footer"
|
||||
|
||||
|
||||
def test_wrong_html_parser_causes_paragraph_to_be_nested_in_div():
|
||||
# This test would fail if html5lib parser would be applied on the input HTML.
|
||||
# It would result in Page: <p></p> <address></address>
|
||||
# instead of Page: <p><address></address></p>
|
||||
|
||||
# language=HTML
|
||||
input_html = """
|
||||
<div class="Page">
|
||||
<p class="NarrativeText">
|
||||
<address class="Address">
|
||||
Mountain View, California
|
||||
</address>
|
||||
</p>
|
||||
</div>
|
||||
"""
|
||||
page = parse_html_to_ontology(input_html)
|
||||
|
||||
assert len(page.children) == 1
|
||||
narrative_text = page.children[0]
|
||||
|
||||
assert len(narrative_text.children) == 1
|
||||
address = narrative_text.children[0]
|
||||
|
||||
assert address.text == "Mountain View, California"
|
||||
|
||||
|
||||
def test_when_class_is_missing_it_can_be_inferred_from_type():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<aside>Some text</aside>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<aside class='Sidebar'><p class='Paragraph'>Some text</p></aside>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_when_class_is_wrong_tag_name_is_overwritten():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<p class='Sidebar'>Some text</p>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<aside class='Sidebar'><p class='Paragraph'>Some text</p></aside>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_when_tag_not_supported_by_ontology_and_wrong_then_consider_them_text():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<newtag class="wrongclass">Some text
|
||||
</newtag>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
# TODO (Pluto): Maybe it should be considered as plain text?
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<span class="UncategorizedText">Some text
|
||||
</span>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_div_are_ignored_when_no_attrs():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<div>
|
||||
<input class="RadioButton" name="health-comparison" type="radio"/>
|
||||
</div>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<input class="RadioButton" name="health-comparison" type="radio"/>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_ids_are_preserved():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<div style="background-color: lightblue" id="important_div">
|
||||
<input class="RadioButton" name="health-comparison" type="radio"/>
|
||||
</div>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<div class="UncategorizedText" style="background-color: lightblue" id="important_div">
|
||||
<input class="RadioButton" name="health-comparison" type="radio"/>
|
||||
</div>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
|
||||
page = ontology.children[0]
|
||||
div_obj = page.children[0]
|
||||
assert div_obj.additional_attributes["id"] == "important_div"
|
||||
|
||||
|
||||
def test_br_is_not_considered_uncategorized_text():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<br/>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<br/>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_text_without_tag_is_marked_as_uncategorized_text_when_there_are_other_elements():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
About the same
|
||||
<input class="RadioButton" name="health-comparison" type="radio"/>
|
||||
Some text
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<p class="Paragraph">
|
||||
About the same
|
||||
</p>
|
||||
<input class="RadioButton" name="health-comparison" type="radio"/>
|
||||
<p class="Paragraph">
|
||||
Some text
|
||||
</p>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_keyword_only_attributes_are_preserved_during_mapping():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<input class="FormFieldValue" type="radio" name="options" value="2" checked>
|
||||
"""
|
||||
) # noqa: E501
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<input class="FormFieldValue" type="radio" name="options" value="2" checked>
|
||||
"""
|
||||
) # noqa: E501
|
||||
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_when_unknown_element_keyword_only_attributes_are_preserved_during_mapping():
|
||||
# <input> can be assigned to multiple classes so it is not clear what it is
|
||||
# thus we assign it to UncategorizedText
|
||||
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<form class="Form">
|
||||
<label class="FormField" for="option1">
|
||||
<input type="radio" name="option1" value="2" checked>
|
||||
<span class="UncategorizedText">
|
||||
Option 1 (Checked)
|
||||
</span>
|
||||
</label>
|
||||
</form>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
# TODO(Pluto): Maybe tag also should be overwritten? Or just leave it as it is?
|
||||
# We classify <input> as UncategorizedText but all the text is preserved
|
||||
# for UnstructuredElement so it make sense now as well
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<form class="Form">
|
||||
<label class="FormField" for="option1">
|
||||
<input class="RadioButton" type="radio" name="option1" value="2" checked />
|
||||
<span class="UncategorizedText">
|
||||
Option 1 (Checked)
|
||||
</span>
|
||||
</label>
|
||||
</form>
|
||||
</div>
|
||||
"""
|
||||
) # noqa: E501
|
||||
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_broken_cell_is_not_raising_error():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<table class="Table">
|
||||
<tbody class="TableBody">
|
||||
<tr class="TableRow">
|
||||
<td class="TableCell>11,442,231</td><td class=" tablecell"="">
|
||||
83.64 GiB
|
||||
</td>
|
||||
<th class="TableCellHeader" rowspan="2">
|
||||
Fair Value
|
||||
</th>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<table class="Table">
|
||||
<tbody>
|
||||
<tr>
|
||||
<td tablecell"="">
|
||||
83.64 GiB
|
||||
</td>
|
||||
<th rowspan="2">
|
||||
Fair Value
|
||||
</th>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_table():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<table class="Table">
|
||||
<tbody class="TableBody">
|
||||
<tr class="TableRow">
|
||||
<td class="TableCell">
|
||||
Fair Value1
|
||||
</td>
|
||||
<th class="TableCellHeader" rowspan="2">
|
||||
Fair Value2
|
||||
</th>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<table class="Table">
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>
|
||||
Fair Value1
|
||||
</td>
|
||||
<th rowspan="2">
|
||||
Fair Value2
|
||||
</th>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_table_and_time():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<table class="Table">
|
||||
<thead class='TableHeader'>
|
||||
<tr class="TableRow">
|
||||
<th class="TableCellHeader" colspan="6">
|
||||
Carrying Value
|
||||
</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody class='TableBody'>
|
||||
<tr class="TableRow">
|
||||
<td class="TableCell" colspan="5">
|
||||
<time class="CalendarDate">
|
||||
June 30, 2023
|
||||
</time>
|
||||
</td>
|
||||
<td class="TableCell">
|
||||
<span class="Currency">
|
||||
$—
|
||||
</span>
|
||||
</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<table class="Table">
|
||||
<thead>
|
||||
<tr>
|
||||
<th colspan="6">
|
||||
Carrying Value
|
||||
</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td colspan="5">
|
||||
<time>
|
||||
June 30, 2023
|
||||
</time>
|
||||
</td>
|
||||
<td>
|
||||
<span>
|
||||
$—
|
||||
</span>
|
||||
</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_malformed_html():
|
||||
# language=HTML
|
||||
input_html = """
|
||||
<!DOCTYPE html PUBLIC "-//W3C//DTD HTML 4.01 Transitional//EN"
|
||||
<html>
|
||||
<head>
|
||||
<title>Super Malformed HTML</title>
|
||||
</head>
|
||||
<body class="Document">
|
||||
<!-- Unclosed comment
|
||||
<div class=>
|
||||
<p>Paragraph with missing closing angle bracket
|
||||
<div>
|
||||
<span>
|
||||
<p>Improperly nested paragraph within a span</span>
|
||||
</p>
|
||||
</div>
|
||||
<script>
|
||||
var x = "Unclosed script tag example;
|
||||
</script>
|
||||
<p>Paragraph with invalid characters: <20> <20> <20></p>
|
||||
</div>
|
||||
</html>
|
||||
"""
|
||||
|
||||
# Such malformed HTML won't be returned by html_partitioning as it uses html5lib parser
|
||||
# to imitate the same behaviour it will be first parsed the same way
|
||||
|
||||
input_html = indent_html(input_html, html_parser="html5lib")
|
||||
|
||||
# Ontology has 1 element and everything inside is just Text
|
||||
# language=HTML
|
||||
expected_html = """
|
||||
<body class="Document">
|
||||
|
||||
<p class="Paragraph">
|
||||
Unclosed comment
|
||||
<div class="">
|
||||
<p>
|
||||
Paragraph with missing closing angle bracket
|
||||
<div>
|
||||
<span>
|
||||
<p>
|
||||
Improperly nested paragraph within a span
|
||||
</p>
|
||||
</span>
|
||||
</div>
|
||||
</p>
|
||||
</div>
|
||||
<script>
|
||||
var x = "Unclosed script tag example;
|
||||
</script>
|
||||
<p>
|
||||
Paragraph with invalid characters: <20> <20> <20>
|
||||
</p>
|
||||
</p>
|
||||
</body>
|
||||
"""
|
||||
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(input_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_text_is_wrapped_inside_layout_element():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
Text
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
# language=HTML
|
||||
expected_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<p class='Paragraph'>Text</p>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
expected_html = indent_html(expected_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
parsed_ontology = indent_html(remove_all_ids(ontology.to_html()))
|
||||
|
||||
assert parsed_ontology == expected_html
|
||||
|
||||
|
||||
def test_text_in_form_field_value():
|
||||
# language=HTML
|
||||
input_html = """
|
||||
<div class="Page">
|
||||
<input class="FormFieldValue" value="Random Input Value"/>
|
||||
</div>
|
||||
"""
|
||||
page = parse_html_to_ontology(input_html)
|
||||
|
||||
assert len(page.children) == 1
|
||||
form_field_value = page.children[0]
|
||||
assert form_field_value.text == ""
|
||||
assert form_field_value.to_text() == "Random Input Value"
|
||||
|
||||
|
||||
def test_text_in_form_field_value_with_null_value():
|
||||
# language=HTML
|
||||
input_html = """
|
||||
<div class="Page">
|
||||
<input class="FormFieldValue" value=""/>
|
||||
</div>
|
||||
"""
|
||||
page = parse_html_to_ontology(input_html)
|
||||
|
||||
assert len(page.children) == 1
|
||||
form_field_value = page.children[0]
|
||||
assert form_field_value.text == ""
|
||||
assert form_field_value.to_text() == ""
|
||||
|
||||
|
||||
def test_to_text_when_form_field():
|
||||
ontology = Page(
|
||||
children=[
|
||||
Form(
|
||||
tag="input",
|
||||
additional_attributes={"value": "Random Input Value"},
|
||||
children=[
|
||||
FormFieldValue(
|
||||
tag="input",
|
||||
additional_attributes={"value": "Random Input Value"},
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
assert ontology.to_text(add_children=True) == "Random Input Value"
|
||||
|
||||
|
||||
def test_recursion_limit_is_limiting_parsing():
|
||||
# language=HTML
|
||||
broken_html = "some text"
|
||||
for i in range(100):
|
||||
broken_html = f"<p class='Paragraph'>{broken_html}</p>"
|
||||
broken_html = _wrap_with_body(broken_html)
|
||||
ontology = parse_html_to_ontology(broken_html)
|
||||
|
||||
iterator = 1
|
||||
last_child = ontology.children[0]
|
||||
while last_child.children:
|
||||
last_child = last_child.children[0]
|
||||
iterator += 1
|
||||
assert last_child.text.startswith('<p class="Paragraph">')
|
||||
assert iterator == RECURSION_LIMIT
|
||||
|
||||
|
||||
def test_get_text_when_recursion_limit_activated():
|
||||
broken_html = "some text"
|
||||
for i in range(100):
|
||||
broken_html = f"<p class='Paragraph'>{broken_html}</p>"
|
||||
broken_html = _wrap_with_body(broken_html)
|
||||
ontology = parse_html_to_ontology(broken_html)
|
||||
|
||||
last_child = ontology.children[0]
|
||||
while last_child.children:
|
||||
last_child = last_child.children[0]
|
||||
|
||||
assert last_child.to_text() == "some text"
|
||||
|
||||
|
||||
def test_uncategorizedtest_has_image_and_no_text():
|
||||
# language=HTML
|
||||
base_html = _wrap_with_body(
|
||||
"""
|
||||
<div class="Page">
|
||||
<div class="UncategorizedText">
|
||||
<img src="https://www.example.com/image.jpg"/>
|
||||
</div>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
base_html = indent_html(base_html)
|
||||
|
||||
ontology: OntologyElement = parse_html_to_ontology(base_html)
|
||||
|
||||
element = ontology.children[0].children[0]
|
||||
assert type(element) is Image
|
||||
assert element.css_class_name == "Image"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("input_type", "expected_class"),
|
||||
[
|
||||
("checkbox", Checkbox),
|
||||
("radio", RadioButton),
|
||||
("text", FormFieldValue), # explicit non-specialised type
|
||||
(None, FormFieldValue), # missing type attribute
|
||||
],
|
||||
)
|
||||
def test_input_tag_type_is_mapped_to_correct_ontology_class(
|
||||
input_type: Optional[str], expected_class: Type[OntologyElement]
|
||||
) -> None:
|
||||
"""Ensure bare <input> tags are classified based on their *type* attribute."""
|
||||
|
||||
type_attr = f' type="{input_type}"' if input_type is not None else ""
|
||||
html_snippet = f'<div class="Page"><input{type_attr} name="field" /></div>'
|
||||
|
||||
page = parse_html_to_ontology(html_snippet)
|
||||
assert len(page.children) == 1
|
||||
element = page.children[0]
|
||||
|
||||
# Validate chosen ontology class and preserved HTML semantics
|
||||
assert isinstance(element, expected_class)
|
||||
assert element.html_tag_name == "input"
|
||||
assert element.css_class_name == expected_class.__name__
|
||||
@@ -0,0 +1,502 @@
|
||||
# End 2 End tests for 15 types
|
||||
|
||||
from unstructured.documents.elements import (
|
||||
Element,
|
||||
ElementMetadata,
|
||||
Header,
|
||||
NarrativeText,
|
||||
Table,
|
||||
Text,
|
||||
)
|
||||
from unstructured.documents.ontology import Address, Paragraph
|
||||
from unstructured.partition.html.html_utils import indent_html
|
||||
from unstructured.partition.html.partition import partition_html
|
||||
from unstructured.partition.html.transformations import (
|
||||
ontology_to_unstructured_elements,
|
||||
parse_html_to_ontology,
|
||||
unstructured_elements_to_ontology,
|
||||
)
|
||||
|
||||
|
||||
def _wrap_in_body_and_page(html_code):
|
||||
return (
|
||||
f'<body class="Document">'
|
||||
f'<div class="Page" data-page-number="1">'
|
||||
f"{html_code}"
|
||||
f"</div>"
|
||||
f"</body>"
|
||||
)
|
||||
|
||||
|
||||
_page_elements = [
|
||||
Text(
|
||||
text="",
|
||||
metadata=ElementMetadata(text_as_html='<div class="Page" data-page-number="1" />'),
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _assert_elements_equal(actual_elements: list[Element], expected_elements: list[Element]):
|
||||
assert len(actual_elements) == len(expected_elements)
|
||||
for actual, expected in zip(actual_elements, expected_elements):
|
||||
assert actual == expected, f"Actual: {actual}, Expected: {expected}"
|
||||
# Not all elements are considered be __eq__ Elements method
|
||||
actual_html = indent_html(actual.metadata.text_as_html, html_parser="html.parser")
|
||||
expected_html = indent_html(expected.metadata.text_as_html, html_parser="html.parser")
|
||||
assert actual_html == expected_html, f"Actual: {actual_html}, Expected: {expected_html}"
|
||||
|
||||
|
||||
def _parse_to_unstructured_elements_and_back_to_html(html_as_str: str):
|
||||
unstructured_elements = partition_html(
|
||||
text=html_as_str, add_img_alt_text=False, html_parser_version="v2", unique_element_ids=True
|
||||
)
|
||||
parsed_ontology = unstructured_elements_to_ontology(unstructured_elements)
|
||||
return unstructured_elements, parsed_ontology
|
||||
|
||||
|
||||
def test_simple_narrative_text_with_id():
|
||||
# language=HTML
|
||||
html_as_str = _wrap_in_body_and_page(
|
||||
"""
|
||||
<p class="NarrativeText">
|
||||
DEALER ONLY
|
||||
</p>
|
||||
"""
|
||||
)
|
||||
|
||||
unstructured_elements, parsed_ontology = _parse_to_unstructured_elements_and_back_to_html(
|
||||
html_as_str
|
||||
)
|
||||
|
||||
expected_html = indent_html(html_as_str, html_parser="html.parser")
|
||||
parsed_html = indent_html(parsed_ontology.to_html(), html_parser="html.parser")
|
||||
|
||||
assert expected_html == parsed_html
|
||||
expected_elements = _page_elements + [
|
||||
NarrativeText(
|
||||
text="DEALER ONLY",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<p class="NarrativeText">DEALER ONLY</p>',
|
||||
),
|
||||
)
|
||||
]
|
||||
|
||||
_assert_elements_equal(unstructured_elements, expected_elements)
|
||||
|
||||
|
||||
def test_input_with_radio_button_checked():
|
||||
# language=HTML
|
||||
html_as_str = _wrap_in_body_and_page(
|
||||
"""
|
||||
<input class="RadioButton" name="health-comparison" type="radio" checked/>
|
||||
"""
|
||||
)
|
||||
|
||||
unstructured_elements, parsed_ontology = _parse_to_unstructured_elements_and_back_to_html(
|
||||
html_as_str
|
||||
)
|
||||
|
||||
expected_html = indent_html(html_as_str, html_parser="html.parser")
|
||||
parsed_html = indent_html(parsed_ontology.to_html(), html_parser="html.parser")
|
||||
|
||||
assert expected_html == parsed_html
|
||||
expected_elements = _page_elements + [
|
||||
Text(
|
||||
text="",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html=(
|
||||
'<input class="RadioButton" name="health-comparison"' 'type="radio" checked />'
|
||||
),
|
||||
),
|
||||
)
|
||||
]
|
||||
_assert_elements_equal(unstructured_elements, expected_elements)
|
||||
|
||||
|
||||
def test_multiple_elements():
|
||||
# language=HTML
|
||||
html_as_str = _wrap_in_body_and_page(
|
||||
"""
|
||||
<p class="Paragraph">
|
||||
About the same
|
||||
</p>
|
||||
<input class="RadioButton" name="health-comparison" type="radio"/>
|
||||
<p class="Paragraph">
|
||||
Some text
|
||||
</p>
|
||||
"""
|
||||
)
|
||||
|
||||
unstructured_elements, parsed_ontology = _parse_to_unstructured_elements_and_back_to_html(
|
||||
html_as_str
|
||||
)
|
||||
|
||||
expected_html = indent_html(html_as_str, html_parser="html.parser")
|
||||
parsed_html = indent_html(parsed_ontology.to_html(), html_parser="html.parser")
|
||||
|
||||
assert expected_html == parsed_html
|
||||
expected_elements = _page_elements + [
|
||||
NarrativeText(
|
||||
text="About the same",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<p class="Paragraph">About the same</p>',
|
||||
),
|
||||
),
|
||||
Text(
|
||||
text="",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<input class="RadioButton" name="health-comparison" type="radio" />',
|
||||
),
|
||||
),
|
||||
NarrativeText(
|
||||
text="Some text",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<p class="Paragraph">Some text</p>',
|
||||
),
|
||||
),
|
||||
]
|
||||
_assert_elements_equal(unstructured_elements, expected_elements)
|
||||
|
||||
|
||||
def test_multiple_pages():
|
||||
# language=HTML
|
||||
html_as_str = """
|
||||
<body class="Document">
|
||||
<div class="Page" data-page-number="1">
|
||||
<p class="Paragraph">
|
||||
Some text
|
||||
</p>
|
||||
</div>
|
||||
<div class="Page" data-page-number="2">
|
||||
<p class="Paragraph">
|
||||
Another text
|
||||
</p>
|
||||
</div>
|
||||
</body>
|
||||
"""
|
||||
|
||||
unstructured_elements, parsed_ontology = _parse_to_unstructured_elements_and_back_to_html(
|
||||
html_as_str
|
||||
)
|
||||
|
||||
expected_html = indent_html(html_as_str, html_parser="html.parser")
|
||||
parsed_html = indent_html(parsed_ontology.to_html(), html_parser="html.parser")
|
||||
|
||||
assert expected_html == parsed_html
|
||||
|
||||
expected_elements = [
|
||||
Text(
|
||||
text="",
|
||||
metadata=ElementMetadata(text_as_html='<div class="Page" data-page-number="1" />'),
|
||||
),
|
||||
NarrativeText(
|
||||
text="Some text",
|
||||
metadata=ElementMetadata(text_as_html='<p class="Paragraph">Some text</p>'),
|
||||
),
|
||||
Text(
|
||||
text="",
|
||||
metadata=ElementMetadata(text_as_html='<div class="Page" data-page-number="2" />'),
|
||||
),
|
||||
NarrativeText(
|
||||
text="Another text",
|
||||
metadata=ElementMetadata(text_as_html='<p class="Paragraph">Another text</p>'),
|
||||
),
|
||||
]
|
||||
_assert_elements_equal(unstructured_elements, expected_elements)
|
||||
|
||||
|
||||
def test_forms():
|
||||
# language=HTML
|
||||
html_as_str = _wrap_in_body_and_page(
|
||||
"""
|
||||
<form class="Form">
|
||||
<label class="FormField" for="option1">
|
||||
<input class="FormFieldValue" type="radio"
|
||||
name="options" value="2" checked>
|
||||
<p class="Paragraph">
|
||||
Option 1 (Checked)
|
||||
</p>
|
||||
</label>
|
||||
</form>
|
||||
"""
|
||||
)
|
||||
|
||||
unstructured_elements, parsed_ontology = _parse_to_unstructured_elements_and_back_to_html(
|
||||
html_as_str
|
||||
)
|
||||
|
||||
expected_html = indent_html(html_as_str, html_parser="html.parser")
|
||||
parsed_html = indent_html(parsed_ontology.to_html(), html_parser="html.parser")
|
||||
assert expected_html == parsed_html
|
||||
expected_elements = _page_elements + [
|
||||
Text(
|
||||
text="2 Option 1 (Checked)",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html=""
|
||||
'<form class="Form">'
|
||||
'<label class="FormField" '
|
||||
'for="option1">'
|
||||
'<input class="FormFieldValue" type="radio" '
|
||||
'name="options" value="2" checked />'
|
||||
'<p class="Paragraph">'
|
||||
"Option 1 (Checked)"
|
||||
"</p></label></form>",
|
||||
),
|
||||
)
|
||||
]
|
||||
_assert_elements_equal(unstructured_elements, expected_elements)
|
||||
|
||||
|
||||
def test_table():
|
||||
# language=HTML
|
||||
html_as_str = _wrap_in_body_and_page(
|
||||
"""
|
||||
<table class="Table">
|
||||
<tbody class="TableBody">
|
||||
<tr class="TableRow">
|
||||
<td class="TableCell">
|
||||
Fair Value1
|
||||
</td>
|
||||
<th class="TableCellHeader" rowspan="2">
|
||||
Fair Value2
|
||||
</th>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
"""
|
||||
)
|
||||
|
||||
unstructured_elements, parsed_ontology = _parse_to_unstructured_elements_and_back_to_html(
|
||||
html_as_str
|
||||
)
|
||||
|
||||
expected_elements = _page_elements + [
|
||||
Table(
|
||||
text="Fair Value1 Fair Value2",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<table class="Table">'
|
||||
"<tbody>"
|
||||
"<tr>"
|
||||
"<td>"
|
||||
"Fair Value1"
|
||||
"</td>"
|
||||
'<th rowspan="2">'
|
||||
"Fair Value2"
|
||||
"</th></tr></tbody></table>",
|
||||
),
|
||||
)
|
||||
]
|
||||
_assert_elements_equal(unstructured_elements, expected_elements)
|
||||
|
||||
|
||||
def test_very_nested_structure_is_preserved():
|
||||
# language=HTML
|
||||
html_as_str = _wrap_in_body_and_page(
|
||||
"""
|
||||
<section class='Section'>
|
||||
<div class='Column'>
|
||||
<header class='Header'>
|
||||
<h1 class='Title'>
|
||||
Title
|
||||
</h1>
|
||||
</header>
|
||||
</div>
|
||||
</section>
|
||||
<div class='Column'>
|
||||
<header class='Header'>
|
||||
Page 1
|
||||
</header>
|
||||
<blockquote class="Quote">
|
||||
<p class="Paragraph">
|
||||
Clever Quote
|
||||
</p>
|
||||
</blockquote>
|
||||
<div class='Footnote'>
|
||||
<span class='UncategorizedText'>
|
||||
Uncategorized footnote text
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
unstructured_elements, parsed_ontology = _parse_to_unstructured_elements_and_back_to_html(
|
||||
html_as_str
|
||||
)
|
||||
expected_html = indent_html(html_as_str, html_parser="html.parser")
|
||||
parsed_html = indent_html(parsed_ontology.to_html(), html_parser="html.parser")
|
||||
|
||||
assert expected_html == parsed_html
|
||||
expected_elements = _page_elements + [
|
||||
Text(
|
||||
text="",
|
||||
metadata=ElementMetadata(text_as_html='<section class="Section" />'),
|
||||
),
|
||||
Text(
|
||||
text="",
|
||||
metadata=ElementMetadata(text_as_html='<div class="Column" />'),
|
||||
),
|
||||
Header(
|
||||
text="Title",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<header class="Header"><h1 class="Title">Title</h1></header>'
|
||||
),
|
||||
),
|
||||
Text(
|
||||
text="",
|
||||
metadata=ElementMetadata(text_as_html='<div class="Column" />'),
|
||||
),
|
||||
Header(
|
||||
text="Page 1",
|
||||
metadata=ElementMetadata(text_as_html='<header class="Header">Page 1</header>'),
|
||||
),
|
||||
NarrativeText(
|
||||
text="Clever Quote",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<blockquote class="Quote">'
|
||||
'<p class="Paragraph">'
|
||||
"Clever Quote"
|
||||
"</p>"
|
||||
"</blockquote>",
|
||||
),
|
||||
),
|
||||
Text(
|
||||
text="Uncategorized footnote text",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<div class="Footnote">'
|
||||
'<span class="UncategorizedText">'
|
||||
"Uncategorized footnote text"
|
||||
"</span>"
|
||||
"</div>",
|
||||
),
|
||||
),
|
||||
]
|
||||
_assert_elements_equal(unstructured_elements, expected_elements)
|
||||
|
||||
|
||||
def test_ordered_list():
|
||||
# language=HTML
|
||||
html_as_str = _wrap_in_body_and_page(
|
||||
"""
|
||||
<ul class="UnorderedList">
|
||||
<li class="ListItem">
|
||||
Item 1
|
||||
</li>
|
||||
<li class="ListItem">
|
||||
Item 2
|
||||
</li>
|
||||
<li class="ListItem">
|
||||
Item 3
|
||||
</li>
|
||||
</ul>
|
||||
"""
|
||||
)
|
||||
|
||||
unstructured_elements, parsed_ontology = _parse_to_unstructured_elements_and_back_to_html(
|
||||
html_as_str
|
||||
)
|
||||
expected_html = indent_html(html_as_str, html_parser="html.parser")
|
||||
parsed_html = indent_html(parsed_ontology.to_html(), html_parser="html.parser")
|
||||
|
||||
assert expected_html == parsed_html
|
||||
expected_elements = _page_elements + [
|
||||
Text(
|
||||
text="Item 1 Item 2 Item 3",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<ul class="UnorderedList">'
|
||||
'<li class="ListItem">'
|
||||
"Item 1"
|
||||
"</li>"
|
||||
'<li class="ListItem">'
|
||||
"Item 2</li>"
|
||||
'<li class="ListItem">'
|
||||
"Item 3"
|
||||
"</li></ul>",
|
||||
),
|
||||
)
|
||||
]
|
||||
_assert_elements_equal(unstructured_elements, expected_elements)
|
||||
|
||||
|
||||
def test_squeezed_elements_are_parsed_back():
|
||||
# language=HTML
|
||||
html_as_str = _wrap_in_body_and_page(
|
||||
"""
|
||||
<p class="NarrativeText">
|
||||
Table of Contents
|
||||
</p>
|
||||
<address class="Address">
|
||||
68 Prince Street Palmdale, CA 93550
|
||||
</address>
|
||||
<a class="Hyperlink">
|
||||
www.google.com
|
||||
</a>
|
||||
"""
|
||||
)
|
||||
|
||||
unstructured_elements, parsed_ontology = _parse_to_unstructured_elements_and_back_to_html(
|
||||
html_as_str
|
||||
)
|
||||
expected_html = indent_html(html_as_str, html_parser="html.parser")
|
||||
parsed_html = indent_html(parsed_ontology.to_html(), html_parser="html.parser")
|
||||
|
||||
assert expected_html == parsed_html
|
||||
expected_elements = _page_elements + [
|
||||
NarrativeText(
|
||||
text="Table of Contents 68 Prince Street Palmdale, CA 93550 www.google.com",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<p class="NarrativeText">Table of Contents</p>'
|
||||
'<address class="Address">'
|
||||
"68 Prince Street Palmdale, CA 93550"
|
||||
"</address>"
|
||||
'<a class="Hyperlink">www.google.com</a>',
|
||||
),
|
||||
)
|
||||
]
|
||||
_assert_elements_equal(unstructured_elements, expected_elements)
|
||||
|
||||
|
||||
def test_inline_elements_are_squeezed_when_text_wrapped_into_paragraphs():
|
||||
# language=HTML
|
||||
base_html = """
|
||||
<div class="Page">
|
||||
About the same
|
||||
<address class="Address">
|
||||
1356 Hornor Avenue Oklahoma
|
||||
</address>
|
||||
Some text
|
||||
</div>
|
||||
"""
|
||||
# Such HTML is transformed into Page: [Pargraph, Address, Paragraph]
|
||||
# We would like it to be parsed to UnstructuredElements as [Page, NarrativeText]
|
||||
|
||||
ontology = parse_html_to_ontology(base_html)
|
||||
|
||||
p1, address, p2 = ontology.children
|
||||
assert isinstance(p1, Paragraph)
|
||||
assert isinstance(address, Address)
|
||||
assert isinstance(p2, Paragraph)
|
||||
|
||||
unstructured_elements = ontology_to_unstructured_elements(ontology)
|
||||
|
||||
assert len(unstructured_elements) == 2
|
||||
assert isinstance(unstructured_elements[0], Text)
|
||||
assert isinstance(unstructured_elements[1], NarrativeText)
|
||||
|
||||
|
||||
def test_alternate_text_from_image_is_passed():
|
||||
# language=HTML
|
||||
input_html = """
|
||||
<div class="Page">
|
||||
<table>
|
||||
<tr>
|
||||
<td rowspan="2">Example image nested in the table:</td>
|
||||
<td rowspan="2"><img src="my-logo.png" alt="ALT TEXT Logo"></td>
|
||||
</tr>
|
||||
</table>
|
||||
</div>add_img_alt_text
|
||||
"""
|
||||
page = parse_html_to_ontology(input_html)
|
||||
unstructured_elements = ontology_to_unstructured_elements(page)
|
||||
assert len(unstructured_elements) == 2
|
||||
assert "ALT TEXT Logo" in unstructured_elements[1].text
|
||||
@@ -0,0 +1,26 @@
|
||||
import pytest
|
||||
|
||||
from unstructured.partition.html.transformations import remove_empty_tags_from_html_content
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("html_content, expected_output"), # noqa PT006
|
||||
[
|
||||
("<div></div>", ""),
|
||||
("<div><p></p></div>", "<div></div>"),
|
||||
("<div><input/></div>", "<div><input/></div>"),
|
||||
("<div><br/></div>", "<div><br/></div>"),
|
||||
('<div><p id="1"></p></div>', '<div><p id="1"></p></div>'),
|
||||
("<div><p>Content</p></div>", "<div><p>Content</p></div>"),
|
||||
("<div><p> </p></div>", "<div></div>"),
|
||||
("<div><p></p><span></span></div>", "<div></div>"),
|
||||
("<div><p>Content</p><span></span></div>", "<div><p>Content</p></div>"),
|
||||
("<div><p>Content</p><span> </span></div>", "<div><p>Content</p></div>"),
|
||||
(
|
||||
"<div><p>Content</p><span>Text</span></div>",
|
||||
"<div><p>Content</p><span>Text</span></div>",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_removes_empty_tags(html_content, expected_output):
|
||||
assert remove_empty_tags_from_html_content(html_content) == expected_output
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,77 @@
|
||||
from unstructured.partition.html import partition_html
|
||||
|
||||
|
||||
def test_alternative_image_text_can_be_included():
|
||||
# language=HTML
|
||||
html = """
|
||||
<div class="Page">
|
||||
<img src="my-logo.png" alt="ALT TEXT Logo"/>
|
||||
</div>
|
||||
"""
|
||||
_, image_to_text_alt_mode = partition_html(
|
||||
text=html,
|
||||
image_alt_mode="to_text",
|
||||
html_parser_version="v2",
|
||||
)
|
||||
assert "ALT TEXT Logo" in image_to_text_alt_mode.text
|
||||
|
||||
_, image_none_alt_mode = partition_html(
|
||||
text=html,
|
||||
image_alt_mode=None,
|
||||
html_parser_version="v2",
|
||||
)
|
||||
assert "ALT TEXT Logo" not in image_none_alt_mode.text
|
||||
|
||||
|
||||
def test_alternative_image_text_can_be_included_when_nested_in_paragraph():
|
||||
# language=HTML
|
||||
html = """
|
||||
<div class="Page">
|
||||
<p class="Paragraph">
|
||||
<img src="my-logo.png" alt="ALT TEXT Logo"/>
|
||||
</p>
|
||||
</div>
|
||||
"""
|
||||
_, paragraph_to_text_alt_mode = partition_html(
|
||||
text=html,
|
||||
image_alt_mode="to_text",
|
||||
html_parser_version="v2",
|
||||
)
|
||||
assert "ALT TEXT Logo" in paragraph_to_text_alt_mode.text
|
||||
|
||||
_, paragraph_none_alt_mode = partition_html(
|
||||
text=html,
|
||||
image_alt_mode=None,
|
||||
html_parser_version="v2",
|
||||
)
|
||||
assert "ALT TEXT Logo" not in paragraph_none_alt_mode.text
|
||||
|
||||
|
||||
def test_attr_and_html_inside_table_cell_is_kept():
|
||||
# language=HTML
|
||||
html = """
|
||||
<div class="Page">
|
||||
<table class="Table">
|
||||
<tbody>
|
||||
<tr>
|
||||
<td colspan="2">
|
||||
Some text
|
||||
</td>
|
||||
<td>
|
||||
<input checked="" class="Checkbox" type="checkbox"/>
|
||||
</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
"""
|
||||
page, table = partition_html(
|
||||
text=html,
|
||||
image_alt_mode="to_text",
|
||||
html_parser_version="v2",
|
||||
)
|
||||
|
||||
assert (
|
||||
'<input checked="" class="Checkbox" type="checkbox"/>' in table.metadata.text_as_html
|
||||
) # class is removed
|
||||
assert 'colspan="2"' in table.metadata.text_as_html
|
||||
@@ -0,0 +1,152 @@
|
||||
from unstructured.documents.elements import ElementMetadata, NarrativeText, Text
|
||||
from unstructured.documents.ontology import Document, Page, Paragraph
|
||||
from unstructured.partition.html.transformations import unstructured_elements_to_ontology
|
||||
|
||||
|
||||
def test_when_first_elements_does_not_have_id():
|
||||
unstructured_elements = [
|
||||
Text(
|
||||
element_id="1",
|
||||
text="",
|
||||
metadata=ElementMetadata(text_as_html='<div class="Page"/>'),
|
||||
),
|
||||
NarrativeText(
|
||||
element_id="2",
|
||||
text="Example text",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<p class="Paragraph"> Example text </p>', parent_id="1"
|
||||
),
|
||||
),
|
||||
]
|
||||
ontology = unstructured_elements_to_ontology(unstructured_elements)
|
||||
|
||||
assert isinstance(ontology, Document)
|
||||
|
||||
assert len(ontology.children) == 1
|
||||
page = ontology.children[0]
|
||||
|
||||
assert isinstance(page, Page)
|
||||
assert len(page.children) == 1
|
||||
paragraph = page.children[0]
|
||||
|
||||
assert isinstance(paragraph, Paragraph)
|
||||
assert paragraph.text == "Example text"
|
||||
|
||||
|
||||
def test_when_two_combined_elements_have_the_same_parent():
|
||||
unstructured_elements = [
|
||||
Text(
|
||||
element_id="1",
|
||||
text="",
|
||||
metadata=ElementMetadata(text_as_html='<div class="Page"/>'),
|
||||
),
|
||||
NarrativeText(
|
||||
element_id="2",
|
||||
text="Example text",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html=(
|
||||
'<p class="Paragraph"> Example text </p>'
|
||||
'<p class="Paragraph"> Example text 2 </p>'
|
||||
),
|
||||
parent_id="1",
|
||||
),
|
||||
),
|
||||
NarrativeText(
|
||||
element_id="3",
|
||||
text="Example text 2",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html=(
|
||||
'<p class="Paragraph"> Example text 3 </p>'
|
||||
'<p class="Paragraph"> Example text 4 </p>'
|
||||
),
|
||||
parent_id="1",
|
||||
),
|
||||
),
|
||||
]
|
||||
ontology = unstructured_elements_to_ontology(unstructured_elements)
|
||||
|
||||
assert isinstance(ontology, Document)
|
||||
|
||||
assert len(ontology.children) == 1
|
||||
page = ontology.children[0]
|
||||
|
||||
assert isinstance(page, Page)
|
||||
assert len(page.children) == 4
|
||||
|
||||
|
||||
def test_element_without_parent_isnt_lost():
|
||||
unstructured_elements = [
|
||||
Text(
|
||||
element_id="1",
|
||||
text="",
|
||||
metadata=ElementMetadata(text_as_html='<div class="Page"/>'),
|
||||
),
|
||||
NarrativeText(
|
||||
element_id="2",
|
||||
text="Example text",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<p class="Paragraph"> Example text </p>', parent_id="1"
|
||||
),
|
||||
),
|
||||
NarrativeText(
|
||||
element_id="3",
|
||||
text="Example text without parent",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<p class="Paragraph"> Example text without parent </p>'
|
||||
),
|
||||
),
|
||||
]
|
||||
ontology = unstructured_elements_to_ontology(unstructured_elements)
|
||||
|
||||
assert isinstance(ontology, Document)
|
||||
assert len(ontology.children) == 2
|
||||
page, paragraph = ontology.children
|
||||
assert isinstance(page, Page)
|
||||
assert len(page.children) == 1
|
||||
assert isinstance(paragraph, Paragraph)
|
||||
assert paragraph.text == "Example text without parent"
|
||||
|
||||
|
||||
def test_multiple_pages_can_be_combined():
|
||||
unstructured_elements = [
|
||||
Text(
|
||||
element_id="1",
|
||||
text="",
|
||||
metadata=ElementMetadata(text_as_html='<div class="Page"/>'),
|
||||
),
|
||||
NarrativeText(
|
||||
element_id="2",
|
||||
text="Example text on page 1",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<p class="Paragraph"> Example text on page 1 </p>', parent_id="1"
|
||||
),
|
||||
),
|
||||
Text(
|
||||
element_id="3",
|
||||
text="",
|
||||
metadata=ElementMetadata(text_as_html='<div class="Page"/>'),
|
||||
),
|
||||
NarrativeText(
|
||||
element_id="4",
|
||||
text="Example text on page 2",
|
||||
metadata=ElementMetadata(
|
||||
text_as_html='<p class="Paragraph"> Example text on page 2 </p>', parent_id="3"
|
||||
),
|
||||
),
|
||||
]
|
||||
ontology = unstructured_elements_to_ontology(unstructured_elements)
|
||||
|
||||
assert isinstance(ontology, Document)
|
||||
assert len(ontology.children) == 2
|
||||
page1 = ontology.children[0]
|
||||
page2 = ontology.children[1]
|
||||
assert isinstance(page1, Page)
|
||||
assert isinstance(page2, Page)
|
||||
assert len(page1.children) == 1
|
||||
assert len(page2.children) == 1
|
||||
paragraph1 = page1.children[0]
|
||||
paragraph2 = page2.children[0]
|
||||
assert isinstance(paragraph1, Paragraph)
|
||||
assert isinstance(paragraph2, Paragraph)
|
||||
assert paragraph1.text == "Example text on page 1"
|
||||
assert paragraph2.text == "Example text on page 2"
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,78 @@
|
||||
import pytest
|
||||
from unstructured_inference.inference.elements import EmbeddedTextRegion
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_embedded_text_regions():
|
||||
return [
|
||||
EmbeddedTextRegion.from_coords(
|
||||
x1=453.00277777777774,
|
||||
y1=317.319341111111,
|
||||
x2=711.5338541666665,
|
||||
y2=358.28571222222206,
|
||||
text="LayoutParser:",
|
||||
),
|
||||
EmbeddedTextRegion.from_coords(
|
||||
x1=726.4778125,
|
||||
y1=317.319341111111,
|
||||
x2=760.3308594444444,
|
||||
y2=357.1698966666667,
|
||||
text="A",
|
||||
),
|
||||
EmbeddedTextRegion.from_coords(
|
||||
x1=775.2748177777777,
|
||||
y1=317.319341111111,
|
||||
x2=917.3579885555555,
|
||||
y2=357.1698966666667,
|
||||
text="Unified",
|
||||
),
|
||||
EmbeddedTextRegion.from_coords(
|
||||
x1=932.3019468888888,
|
||||
y1=317.319341111111,
|
||||
x2=1071.8426522222221,
|
||||
y2=357.1698966666667,
|
||||
text="Toolkit",
|
||||
),
|
||||
EmbeddedTextRegion.from_coords(
|
||||
x1=1086.7866105555556,
|
||||
y1=317.319341111111,
|
||||
x2=1141.2105142777777,
|
||||
y2=357.1698966666667,
|
||||
text="for",
|
||||
),
|
||||
EmbeddedTextRegion.from_coords(
|
||||
x1=1156.154472611111,
|
||||
y1=317.319341111111,
|
||||
x2=1256.334784222222,
|
||||
y2=357.1698966666667,
|
||||
text="Deep",
|
||||
),
|
||||
EmbeddedTextRegion.from_coords(
|
||||
x1=437.83888888888885,
|
||||
y1=367.13322999999986,
|
||||
x2=610.0171992222222,
|
||||
y2=406.9837855555556,
|
||||
text="Learning",
|
||||
),
|
||||
EmbeddedTextRegion.from_coords(
|
||||
x1=624.9611575555555,
|
||||
y1=367.13322999999986,
|
||||
x2=741.6754646666665,
|
||||
y2=406.9837855555556,
|
||||
text="Based",
|
||||
),
|
||||
EmbeddedTextRegion.from_coords(
|
||||
x1=756.619423,
|
||||
y1=367.13322999999986,
|
||||
x2=958.3867708333332,
|
||||
y2=406.9837855555556,
|
||||
text="Document",
|
||||
),
|
||||
EmbeddedTextRegion.from_coords(
|
||||
x1=973.3307291666665,
|
||||
y1=367.13322999999986,
|
||||
x2=1092.0535042777776,
|
||||
y2=406.9837855555556,
|
||||
text="Image",
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,154 @@
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
from unstructured_inference.inference.elements import Rectangle
|
||||
from unstructured_inference.inference.layout import DocumentLayout, PageLayout
|
||||
from unstructured_inference.inference.layoutelement import LayoutElement
|
||||
|
||||
from unstructured.partition.pdf_image.analysis.bbox_visualisation import (
|
||||
TextAlignment,
|
||||
get_bbox_text_size,
|
||||
get_bbox_thickness,
|
||||
get_label_rect_and_coords,
|
||||
get_rgb_color,
|
||||
get_text_color,
|
||||
)
|
||||
from unstructured.partition.pdf_image.analysis.layout_dump import ObjectDetectionLayoutDumper
|
||||
|
||||
|
||||
@pytest.mark.parametrize("color", ["red", "green", "blue", "yellow", "black", "white"])
|
||||
def test_get_rgb_color(color: str):
|
||||
color_tuple = get_rgb_color(color)
|
||||
|
||||
assert isinstance(color_tuple, tuple)
|
||||
assert len(color_tuple) == 3
|
||||
assert all(isinstance(c, int) for c in color_tuple)
|
||||
assert all(0 <= c <= 255 for c in color_tuple)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("bbox", "expected_text_size"),
|
||||
[
|
||||
((0, 0, 90, 90), 17),
|
||||
((0, 0, 500, 200), 21),
|
||||
((0, 0, 10000, 10000), 32),
|
||||
],
|
||||
)
|
||||
def test_get_bbox_text_size(bbox: tuple[int, int, int, int], expected_text_size):
|
||||
page_size = (1700, 2200) # standard size of a page
|
||||
text_size = get_bbox_text_size(bbox, page_size)
|
||||
|
||||
assert text_size == expected_text_size
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("bbox", "expected_box_thickness"),
|
||||
[
|
||||
((0, 0, 90, 90), 1),
|
||||
((0, 0, 450, 250), 2),
|
||||
((0, 0, 600, 1000), 3),
|
||||
],
|
||||
)
|
||||
def test_get_bbox_thickness(bbox: tuple[int, int, int, int], expected_box_thickness):
|
||||
page_size = (1700, 2200) # standard size of a page
|
||||
box_thickness = get_bbox_thickness(bbox, page_size)
|
||||
|
||||
assert box_thickness == expected_box_thickness
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("color", "expected_text_color"),
|
||||
[
|
||||
("navy", "white"),
|
||||
("crimson", "white"),
|
||||
("maroon", "white"),
|
||||
("dimgray", "white"),
|
||||
("darkgreen", "white"),
|
||||
("darkcyan", "white"),
|
||||
("fuchsia", "white"),
|
||||
("violet", "black"),
|
||||
("gold", "black"),
|
||||
("aqua", "black"),
|
||||
("greenyellow", "black"),
|
||||
],
|
||||
)
|
||||
def test_best_text_color(color, expected_text_color):
|
||||
color_tuple = get_rgb_color(color)
|
||||
expected_text_color_tuple = get_rgb_color(expected_text_color)
|
||||
|
||||
_, text_color_tuple = get_text_color(color_tuple)
|
||||
assert text_color_tuple == expected_text_color_tuple
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("alignment", "expected_text_bbox"),
|
||||
[
|
||||
(TextAlignment.CENTER, ((145, 145), (155, 155))),
|
||||
(TextAlignment.TOP_LEFT, ((100, 90), (120, 100))),
|
||||
(TextAlignment.TOP_RIGHT, ((180, 100), (200, 110))),
|
||||
(TextAlignment.BOTTOM_LEFT, ((100, 190), (120, 200))),
|
||||
(TextAlignment.BOTTOM_RIGHT, ((180, 190), (200, 200))),
|
||||
],
|
||||
)
|
||||
def test_get_text_bbox(alignment, expected_text_bbox):
|
||||
text_bbox, text_xy = get_label_rect_and_coords(
|
||||
alignment=alignment, bbox_points=(100, 100, 200, 200), text_width=10, text_height=10
|
||||
)
|
||||
# adding high atol to account for the text-based extending of the bbox
|
||||
assert np.allclose(text_bbox, expected_text_bbox, atol=10)
|
||||
|
||||
|
||||
def test_od_document_layout_dump():
|
||||
page1 = PageLayout(
|
||||
number=1,
|
||||
image=Image.new("1", (1, 1)),
|
||||
image_metadata={"width": 100, "height": 100},
|
||||
)
|
||||
page1.elements = [
|
||||
LayoutElement(type="Title", bbox=Rectangle(x1=0, y1=0, x2=10, y2=10), prob=0.7),
|
||||
LayoutElement(type="Paragraph", bbox=Rectangle(x1=0, y1=100, x2=10, y2=110), prob=0.8),
|
||||
]
|
||||
page2 = PageLayout(
|
||||
number=2,
|
||||
image=Image.new("1", (1, 1)),
|
||||
image_metadata={"width": 100, "height": 100},
|
||||
)
|
||||
page2.elements = [
|
||||
LayoutElement(type="Table", bbox=Rectangle(x1=0, y1=0, x2=10, y2=10), prob=0.9),
|
||||
LayoutElement(type="Image", bbox=Rectangle(x1=0, y1=100, x2=10, y2=110), prob=1.0),
|
||||
]
|
||||
od_document_layout = DocumentLayout(pages=[page1, page2])
|
||||
|
||||
expected_dump = {
|
||||
"pages": [
|
||||
{
|
||||
"number": 1,
|
||||
"size": {
|
||||
"width": 100,
|
||||
"height": 100,
|
||||
},
|
||||
"elements": [
|
||||
{"bbox": [0, 0, 10, 10], "type": "Title", "prob": 0.7},
|
||||
{"bbox": [0, 100, 10, 110], "type": "Paragraph", "prob": 0.8},
|
||||
],
|
||||
},
|
||||
{
|
||||
"number": 2,
|
||||
"size": {
|
||||
"width": 100,
|
||||
"height": 100,
|
||||
},
|
||||
"elements": [
|
||||
{"bbox": [0, 0, 10, 10], "type": "Table", "prob": 0.9},
|
||||
{"bbox": [0, 100, 10, 110], "type": "Image", "prob": 1.0},
|
||||
],
|
||||
},
|
||||
]
|
||||
}
|
||||
od_layout_dump = ObjectDetectionLayoutDumper(od_document_layout).dump()
|
||||
|
||||
assert expected_dump == {"pages": od_layout_dump.get("pages")}
|
||||
|
||||
# check OD model classes are attached but do not depend on a specific model instance
|
||||
assert "object_detection_classes" in od_layout_dump
|
||||
assert len(od_layout_dump["object_detection_classes"]) > 0
|
||||
@@ -0,0 +1,670 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
import tempfile
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
from pytest_mock import MockFixture
|
||||
from unstructured_inference.inference import layout
|
||||
from unstructured_pytesseract import TesseractError
|
||||
|
||||
from test_unstructured.partition.pdf_image.test_pdf import assert_element_extraction
|
||||
from test_unstructured.unit_utils import assert_round_trips_through_JSON, example_doc_path
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import ElementType
|
||||
from unstructured.partition import image, pdf
|
||||
from unstructured.partition.pdf_image import ocr
|
||||
from unstructured.partition.utils.constants import (
|
||||
UNSTRUCTURED_INCLUDE_DEBUG_METADATA,
|
||||
PartitionStrategy,
|
||||
)
|
||||
from unstructured.utils import only
|
||||
|
||||
DIRECTORY = pathlib.Path(__file__).parent.resolve()
|
||||
|
||||
|
||||
class MockResponse:
|
||||
def __init__(self, status_code, response):
|
||||
self.status_code = status_code
|
||||
self.response = response
|
||||
|
||||
def json(self):
|
||||
return self.response
|
||||
|
||||
|
||||
def mock_healthy_get(url, **kwargs):
|
||||
return MockResponse(status_code=200, response={})
|
||||
|
||||
|
||||
def mock_unhealthy_get(url, **kwargs):
|
||||
return MockResponse(status_code=500, response={})
|
||||
|
||||
|
||||
def mock_unsuccessful_post(url, **kwargs):
|
||||
return MockResponse(status_code=500, response={})
|
||||
|
||||
|
||||
def mock_successful_post(url, **kwargs):
|
||||
response = {
|
||||
"pages": [
|
||||
{
|
||||
"number": 0,
|
||||
"elements": [
|
||||
{"type": "Title", "text": "Charlie Brown and the Great Pumpkin"},
|
||||
],
|
||||
},
|
||||
{
|
||||
"number": 1,
|
||||
"elements": [{"type": "Title", "text": "A Charlie Brown Christmas"}],
|
||||
},
|
||||
],
|
||||
}
|
||||
return MockResponse(status_code=200, response=response)
|
||||
|
||||
|
||||
class MockPageLayout(layout.PageLayout):
|
||||
def __init__(self, number: int, image: Image):
|
||||
self.number = number
|
||||
self.image = image
|
||||
self.elements = [
|
||||
layout.LayoutElement.from_coords(
|
||||
type="Title",
|
||||
x1=0,
|
||||
y1=0,
|
||||
x2=2,
|
||||
y2=2,
|
||||
text="Charlie Brown and the Great Pumpkin",
|
||||
),
|
||||
]
|
||||
self.elements_array = layout.LayoutElements.from_list(self.elements)
|
||||
|
||||
|
||||
class MockDocumentLayout(layout.DocumentLayout):
|
||||
@property
|
||||
def pages(self):
|
||||
return [
|
||||
MockPageLayout(number=0, image=Image.new("1", (1, 1))),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "file"),
|
||||
[
|
||||
(example_doc_path("img/example.jpg"), None),
|
||||
(None, b"0000"),
|
||||
],
|
||||
)
|
||||
def test_partition_image_local(monkeypatch, filename, file):
|
||||
monkeypatch.setattr(
|
||||
layout,
|
||||
"process_data_with_model",
|
||||
lambda *args, **kwargs: MockDocumentLayout(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
layout,
|
||||
"process_file_with_model",
|
||||
lambda *args, **kwargs: MockDocumentLayout(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ocr,
|
||||
"process_data_with_ocr",
|
||||
lambda *args, **kwargs: MockDocumentLayout(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ocr,
|
||||
"process_data_with_ocr",
|
||||
lambda *args, **kwargs: MockDocumentLayout(),
|
||||
)
|
||||
|
||||
partition_image_response = pdf._partition_pdf_or_image_local(
|
||||
filename,
|
||||
file,
|
||||
is_image=True,
|
||||
)
|
||||
assert partition_image_response[0].text == "Charlie Brown and the Great Pumpkin"
|
||||
|
||||
|
||||
@pytest.mark.skip("Needs to be fixed upstream in unstructured-inference")
|
||||
def test_partition_image_local_raises_with_no_filename():
|
||||
with pytest.raises(FileNotFoundError):
|
||||
pdf._partition_pdf_or_image_local(filename="", file=None, is_image=True)
|
||||
|
||||
|
||||
def test_partition_image_with_auto_strategy():
|
||||
filename = example_doc_path("img/layout-parser-paper-fast.jpg")
|
||||
elements = image.partition_image(filename=filename, strategy=PartitionStrategy.AUTO)
|
||||
titles = [
|
||||
el for el in elements if el.category == ElementType.TITLE and len(el.text.split(" ")) > 10
|
||||
]
|
||||
title = "LayoutParser: A Unified Toolkit for Deep Learning Based Document Image Analysis"
|
||||
idx = 3
|
||||
assert titles[0].text == title
|
||||
assert elements[idx].metadata.detection_class_prob is not None
|
||||
assert isinstance(elements[idx].metadata.detection_class_prob, float)
|
||||
|
||||
|
||||
def test_partition_image_with_table_extraction():
|
||||
filename = example_doc_path("img/layout-parser-paper-with-table.jpg")
|
||||
elements = image.partition_image(
|
||||
filename=filename,
|
||||
strategy=PartitionStrategy.HI_RES,
|
||||
infer_table_structure=True,
|
||||
)
|
||||
table = [el.metadata.text_as_html for el in elements if el.metadata.text_as_html]
|
||||
assert len(table) == 1
|
||||
assert "<table><thead><tr>" in table[0]
|
||||
assert "</thead><tbody><tr>" in table[0]
|
||||
|
||||
|
||||
def test_partition_image_with_multipage_tiff():
|
||||
filename = example_doc_path("img/layout-parser-paper-combined.tiff")
|
||||
elements = image.partition_image(filename=filename, strategy=PartitionStrategy.AUTO)
|
||||
assert elements[-1].metadata.page_number == 2
|
||||
|
||||
|
||||
def test_partition_image_with_bmp(tmpdir):
|
||||
filename = example_doc_path("img/layout-parser-paper-with-table.jpg")
|
||||
bmp_filename = os.path.join(tmpdir.dirname, "example.bmp")
|
||||
img = Image.open(filename)
|
||||
img.save(bmp_filename)
|
||||
|
||||
elements = image.partition_image(
|
||||
filename=bmp_filename,
|
||||
strategy=PartitionStrategy.HI_RES,
|
||||
infer_table_structure=True,
|
||||
)
|
||||
table = [el.metadata.text_as_html for el in elements if el.metadata.text_as_html]
|
||||
assert len(table) == 1
|
||||
assert "<table><thead><tr>" in table[0]
|
||||
assert "</thead><tbody><tr>" in table[0]
|
||||
|
||||
|
||||
def test_partition_image_with_language_passed():
|
||||
filename = example_doc_path("img/example.jpg")
|
||||
with mock.patch.object(
|
||||
ocr,
|
||||
"process_file_with_ocr",
|
||||
mock.MagicMock(),
|
||||
) as mock_partition:
|
||||
image.partition_image(
|
||||
filename=filename,
|
||||
strategy=PartitionStrategy.HI_RES,
|
||||
ocr_languages="eng+swe",
|
||||
)
|
||||
|
||||
assert mock_partition.call_args.kwargs.get("ocr_languages") == "eng+swe"
|
||||
|
||||
|
||||
def test_partition_image_from_file_with_language_passed():
|
||||
filename = example_doc_path("img/example.jpg")
|
||||
with mock.patch.object(
|
||||
ocr,
|
||||
"process_data_with_ocr",
|
||||
mock.MagicMock(),
|
||||
) as mock_partition, open(filename, "rb") as f:
|
||||
image.partition_image(file=f, strategy=PartitionStrategy.HI_RES, ocr_languages="eng+swe")
|
||||
|
||||
assert mock_partition.call_args.kwargs.get("ocr_languages") == "eng+swe"
|
||||
|
||||
|
||||
# NOTE(crag): see https://github.com/Unstructured-IO/unstructured/issues/1086
|
||||
@pytest.mark.skip(reason="Current catching too many tesseract errors")
|
||||
def test_partition_image_raises_with_invalid_language():
|
||||
filename = example_doc_path("img/example.jpg")
|
||||
with pytest.raises(TesseractError):
|
||||
image.partition_image(
|
||||
filename=filename,
|
||||
strategy=PartitionStrategy.HI_RES,
|
||||
ocr_languages="fakeroo",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"strategy",
|
||||
[
|
||||
PartitionStrategy.HI_RES,
|
||||
PartitionStrategy.OCR_ONLY,
|
||||
],
|
||||
)
|
||||
def test_partition_image_strategies_keep_languages_metadata(strategy):
|
||||
filename = example_doc_path("img/english-and-korean.png")
|
||||
elements = image.partition_image(
|
||||
filename=filename,
|
||||
languages=["eng", "kor"],
|
||||
strategy=strategy,
|
||||
)
|
||||
|
||||
assert elements[0].metadata.languages == ["eng", "kor"]
|
||||
|
||||
|
||||
def test_partition_image_with_ocr_detects_korean():
|
||||
filename = example_doc_path("img/english-and-korean.png")
|
||||
elements = image.partition_image(
|
||||
filename=filename,
|
||||
ocr_languages="eng+kor",
|
||||
strategy=PartitionStrategy.OCR_ONLY,
|
||||
)
|
||||
|
||||
assert elements[0].text == "RULES AND INSTRUCTIONS"
|
||||
# FIXME (yao): revisit this lstrip after refactoring merging logics; right now on docker and
|
||||
# local testing yield different results and on docker there is a "," at the start of the Korean
|
||||
# text line
|
||||
assert elements[3].text.replace(" ", "").lstrip(",").startswith("안녕하세요")
|
||||
|
||||
|
||||
def test_partition_image_with_ocr_detects_korean_from_file():
|
||||
filename = example_doc_path("img/english-and-korean.png")
|
||||
with open(filename, "rb") as f:
|
||||
elements = image.partition_image(
|
||||
file=f,
|
||||
ocr_languages="eng+kor",
|
||||
strategy=PartitionStrategy.OCR_ONLY,
|
||||
)
|
||||
|
||||
assert elements[0].text == "RULES AND INSTRUCTIONS"
|
||||
assert elements[3].text.replace(" ", "").lstrip(",").startswith("안녕하세요")
|
||||
|
||||
|
||||
def test_partition_image_raises_with_bad_strategy():
|
||||
filename = example_doc_path("img/english-and-korean.png")
|
||||
with pytest.raises(ValueError):
|
||||
image.partition_image(filename=filename, strategy="fakeroo")
|
||||
|
||||
|
||||
def test_partition_image_default_strategy_hi_res():
|
||||
filename = example_doc_path("img/layout-parser-paper-fast.jpg")
|
||||
with open(filename, "rb") as f:
|
||||
elements = image.partition_image(file=f)
|
||||
|
||||
title = "LayoutParser: A Unified Toolkit for Deep Learning Based Document Image Analysis"
|
||||
idx = 2
|
||||
assert elements[idx].text == title
|
||||
assert elements[idx].metadata.coordinates is not None
|
||||
assert elements[idx].metadata.detection_class_prob is not None
|
||||
assert isinstance(elements[idx].metadata.detection_class_prob, float)
|
||||
if UNSTRUCTURED_INCLUDE_DEBUG_METADATA:
|
||||
# A bug in partition_groups_from_regions in unstructured-inference losses some sources
|
||||
assert {element.metadata.detection_origin for element in elements} == {
|
||||
"yolox",
|
||||
"ocr_tesseract",
|
||||
}
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_image_from_file_path_gets_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.pdf.get_last_modified_date",
|
||||
return_value=filesystem_last_modified,
|
||||
)
|
||||
|
||||
elements = image.partition_image(example_doc_path("img/english-and-korean.png"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_image_from_file_path_with_hi_res_strategy_gets_last_modified_from_filesystem(
|
||||
mocker: MockFixture,
|
||||
):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.pdf.get_last_modified_date",
|
||||
return_value=filesystem_last_modified,
|
||||
)
|
||||
|
||||
elements = image.partition_image(
|
||||
example_doc_path("img/english-and-korean.png"), strategy=PartitionStrategy.HI_RES
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_image_from_file_path_prefers_metadata_last_modified(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
metadata_last_modified = "2009-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.pdf.get_last_modified_date",
|
||||
return_value=filesystem_last_modified,
|
||||
)
|
||||
|
||||
elements = image.partition_image(
|
||||
example_doc_path("img/english-and-korean.png"),
|
||||
metadata_last_modified=metadata_last_modified,
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_image_from_file_path_with_hi_res_strategy_prefers_metadata_last_modified(
|
||||
mocker: MockFixture,
|
||||
):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
metadata_last_modified = "2009-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.pdf.get_last_modified_date",
|
||||
return_value=filesystem_last_modified,
|
||||
)
|
||||
|
||||
elements = image.partition_image(
|
||||
example_doc_path("img/english-and-korean.png"),
|
||||
strategy=PartitionStrategy.HI_RES,
|
||||
metadata_last_modified=metadata_last_modified,
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_image_from_file_gets_last_modified_None():
|
||||
with open(example_doc_path("img/english-and-korean.png"), "rb") as f:
|
||||
elements = image.partition_image(file=f)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_image_from_file_with_hi_res_strategy_gets_last_modified_None(
|
||||
mocker: MockFixture,
|
||||
):
|
||||
with open(example_doc_path("img/english-and-korean.png"), "rb") as f:
|
||||
elements = image.partition_image(file=f, strategy=PartitionStrategy.HI_RES)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_image_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2009-07-05T09:24:28"
|
||||
|
||||
with open(example_doc_path("img/english-and-korean.png"), "rb") as f:
|
||||
elements = image.partition_image(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_image_from_file_with_hi_res_strategy_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2009-07-05T09:24:28"
|
||||
|
||||
with open(example_doc_path("img/english-and-korean.png"), "rb") as f:
|
||||
elements = image.partition_image(
|
||||
file=f,
|
||||
metadata_last_modified=metadata_last_modified,
|
||||
strategy=PartitionStrategy.HI_RES,
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_msg_with_json():
|
||||
elements = image.partition_image(
|
||||
example_doc_path("img/layout-parser-paper-fast.jpg"),
|
||||
strategy=PartitionStrategy.AUTO,
|
||||
)
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
def test_partition_image_with_ocr_has_coordinates_from_filename():
|
||||
filename = example_doc_path("img/english-and-korean.png")
|
||||
elements = image.partition_image(filename=filename, strategy=PartitionStrategy.OCR_ONLY)
|
||||
int_coordinates = [(int(x), int(y)) for x, y in elements[0].metadata.coordinates.points]
|
||||
assert int_coordinates == [(14, 16), (14, 37), (381, 37), (381, 16)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename",
|
||||
[
|
||||
"img/layout-parser-paper-with-table.jpg",
|
||||
"img/english-and-korean.png",
|
||||
"img/layout-parser-paper-fast.jpg",
|
||||
],
|
||||
)
|
||||
def test_partition_image_with_ocr_coordinates_are_not_nan_from_filename(
|
||||
filename,
|
||||
):
|
||||
import math
|
||||
|
||||
elements = image.partition_image(
|
||||
filename=example_doc_path(filename), strategy=PartitionStrategy.OCR_ONLY
|
||||
)
|
||||
for element in elements:
|
||||
# TODO (jennings) One or multiple elements is an empty string
|
||||
# without coordinates. This should be fixed in a new issue
|
||||
if element.text:
|
||||
box = element.metadata.coordinates.points
|
||||
for point in box:
|
||||
assert point[0] is not math.nan
|
||||
assert point[1] is not math.nan
|
||||
|
||||
|
||||
def test_partition_image_formats_languages_for_tesseract():
|
||||
filename = example_doc_path("img/jpn-vert.jpeg")
|
||||
with mock.patch(
|
||||
"unstructured.partition.pdf_image.ocr.process_file_with_ocr",
|
||||
) as mock_process_file_with_ocr:
|
||||
image.partition_image(
|
||||
filename=filename, strategy=PartitionStrategy.HI_RES, languages=["jpn_vert"]
|
||||
)
|
||||
_, kwargs = mock_process_file_with_ocr.call_args_list[0]
|
||||
assert "ocr_languages" in kwargs
|
||||
assert kwargs["ocr_languages"] == "jpn_vert"
|
||||
|
||||
|
||||
def test_partition_image_warns_with_ocr_languages(caplog):
|
||||
filename = example_doc_path("img/layout-parser-paper-fast.jpg")
|
||||
image.partition_image(filename=filename, strategy=PartitionStrategy.HI_RES, ocr_languages="eng")
|
||||
assert "The ocr_languages kwarg will be deprecated" in caplog.text
|
||||
|
||||
|
||||
def test_add_chunking_strategy_on_partition_image():
|
||||
filename = example_doc_path("img/layout-parser-paper-fast.jpg")
|
||||
elements = image.partition_image(filename=filename)
|
||||
chunk_elements = image.partition_image(filename, chunking_strategy="by_title")
|
||||
chunks = chunk_by_title(elements)
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_add_chunking_strategy_on_partition_image_hi_res():
|
||||
filename = example_doc_path("img/layout-parser-paper-with-table.jpg")
|
||||
elements = image.partition_image(
|
||||
filename=filename,
|
||||
strategy=PartitionStrategy.HI_RES,
|
||||
infer_table_structure=True,
|
||||
)
|
||||
chunk_elements = image.partition_image(
|
||||
filename,
|
||||
strategy=PartitionStrategy.HI_RES,
|
||||
infer_table_structure=True,
|
||||
chunking_strategy="by_title",
|
||||
)
|
||||
chunks = chunk_by_title(elements)
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_partition_image_uses_model_name():
|
||||
with mock.patch.object(
|
||||
pdf,
|
||||
"_partition_pdf_or_image_local",
|
||||
) as mockpartition:
|
||||
image.partition_image(
|
||||
example_doc_path("img/layout-parser-paper-fast.jpg"), model_name="test"
|
||||
)
|
||||
print(mockpartition.call_args)
|
||||
assert "model_name" in mockpartition.call_args.kwargs
|
||||
assert mockpartition.call_args.kwargs["model_name"]
|
||||
|
||||
|
||||
def test_partition_image_uses_hi_res_model_name():
|
||||
with mock.patch.object(
|
||||
pdf,
|
||||
"_partition_pdf_or_image_local",
|
||||
) as mockpartition:
|
||||
image.partition_image(
|
||||
example_doc_path("img/layout-parser-paper-fast.jpg"), hi_res_model_name="test"
|
||||
)
|
||||
print(mockpartition.call_args)
|
||||
assert "model_name" not in mockpartition.call_args.kwargs
|
||||
assert "hi_res_model_name" in mockpartition.call_args.kwargs
|
||||
assert mockpartition.call_args.kwargs["hi_res_model_name"] == "test"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("ocr_mode", "idx_title_element"),
|
||||
[
|
||||
("entire_page", 2),
|
||||
("individual_blocks", 1),
|
||||
],
|
||||
)
|
||||
def test_partition_image_hi_res_ocr_mode(ocr_mode, idx_title_element):
|
||||
filename = example_doc_path("img/layout-parser-paper-fast.jpg")
|
||||
elements = image.partition_image(
|
||||
filename=filename, ocr_mode=ocr_mode, strategy=PartitionStrategy.HI_RES
|
||||
)
|
||||
# Note(yuming): idx_title_element is different based on xy-cut and ocr mode
|
||||
assert elements[idx_title_element].category == ElementType.TITLE
|
||||
|
||||
|
||||
def test_partition_image_hi_res_invalid_ocr_mode():
|
||||
filename = example_doc_path("img/layout-parser-paper-fast.jpg")
|
||||
with pytest.raises(ValueError):
|
||||
_ = image.partition_image(
|
||||
filename=filename, ocr_mode="invalid_ocr_mode", strategy=PartitionStrategy.HI_RES
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ocr_mode",
|
||||
[
|
||||
"entire_page",
|
||||
"individual_blocks",
|
||||
],
|
||||
)
|
||||
def test_partition_image_hi_res_ocr_mode_with_table_extraction(ocr_mode):
|
||||
filename = example_doc_path("img/layout-parser-paper-with-table.jpg")
|
||||
elements = image.partition_image(
|
||||
filename=filename,
|
||||
ocr_mode=ocr_mode,
|
||||
strategy=PartitionStrategy.HI_RES,
|
||||
infer_table_structure=True,
|
||||
)
|
||||
table = [el.metadata.text_as_html for el in elements if el.metadata.text_as_html]
|
||||
assert len(table) == 1
|
||||
assert "<table><thead><tr>" in table[0]
|
||||
assert "</thead><tbody><tr>" in table[0]
|
||||
assert "Layouts of history Japanese documents" in table[0]
|
||||
assert "Layouts of scanned modern magazines and scientific reports" in table[0]
|
||||
|
||||
|
||||
def test_partition_image_raises_type_error_for_invalid_languages():
|
||||
filename = example_doc_path("img/layout-parser-paper-fast.jpg")
|
||||
with pytest.raises(TypeError):
|
||||
image.partition_image(filename=filename, strategy=PartitionStrategy.HI_RES, languages="eng")
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def inference_results():
|
||||
page = layout.PageLayout(
|
||||
number=1,
|
||||
image=mock.MagicMock(format="JPEG"),
|
||||
)
|
||||
page.elements = [layout.LayoutElement.from_coords(0, 0, 600, 800, text="hello")]
|
||||
page.elements_array = layout.LayoutElements.from_list(page.elements)
|
||||
doc = layout.DocumentLayout(pages=[page])
|
||||
return doc
|
||||
|
||||
|
||||
def test_partition_image_has_filename(inference_results):
|
||||
filename = "layout-parser-paper-fast.jpg"
|
||||
# Mock inference call with known return results
|
||||
with mock.patch(
|
||||
"unstructured_inference.inference.layout.process_file_with_model",
|
||||
return_value=inference_results,
|
||||
) as mock_inference_func:
|
||||
elements = image.partition_image(
|
||||
filename=example_doc_path(f"img/{filename}"),
|
||||
strategy=PartitionStrategy.HI_RES,
|
||||
)
|
||||
# Make sure we actually went down the path we expect.
|
||||
mock_inference_func.assert_called_once()
|
||||
# Unpack element but also make sure there is only one
|
||||
element = only(elements)
|
||||
# This makes sure we are still getting the filetype metadata (should be translated from the
|
||||
# fixtures)
|
||||
assert element.metadata.filetype == "JPEG"
|
||||
# This should be kept from the filename we originally gave
|
||||
assert element.metadata.filename == filename
|
||||
|
||||
|
||||
@pytest.mark.parametrize("file_mode", ["filename", "rb"])
|
||||
@pytest.mark.parametrize("extract_image_block_to_payload", [False, True])
|
||||
def test_partition_image_element_extraction(
|
||||
file_mode,
|
||||
extract_image_block_to_payload,
|
||||
):
|
||||
filename = example_doc_path("img/embedded-images-tables.jpg")
|
||||
extract_image_block_types = ["Image", "Table"]
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
if file_mode == "filename":
|
||||
elements = image.partition_image(
|
||||
filename=filename,
|
||||
extract_image_block_types=extract_image_block_types,
|
||||
extract_image_block_to_payload=extract_image_block_to_payload,
|
||||
extract_image_block_output_dir=tmpdir,
|
||||
)
|
||||
else:
|
||||
with open(filename, "rb") as f:
|
||||
elements = image.partition_image(
|
||||
file=f,
|
||||
extract_image_block_types=extract_image_block_types,
|
||||
extract_image_block_to_payload=extract_image_block_to_payload,
|
||||
extract_image_block_output_dir=tmpdir,
|
||||
)
|
||||
|
||||
assert_element_extraction(
|
||||
elements, extract_image_block_types, extract_image_block_to_payload, tmpdir
|
||||
)
|
||||
|
||||
|
||||
def test_partition_image_works_on_heic_file():
|
||||
filename = example_doc_path("img/DA-1p.heic")
|
||||
elements = image.partition_image(filename=filename, strategy=PartitionStrategy.AUTO)
|
||||
titles = [el.text for el in elements if el.category == ElementType.TITLE]
|
||||
assert "CREATURES" in titles
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"strategy",
|
||||
[PartitionStrategy.HI_RES, PartitionStrategy.OCR_ONLY],
|
||||
)
|
||||
def test_deterministic_element_ids(strategy: str):
|
||||
elements_1 = image.partition_image(
|
||||
example_doc_path("img/layout-parser-paper-with-table.jpg"),
|
||||
strategy=strategy,
|
||||
starting_page_number=2,
|
||||
)
|
||||
elements_2 = image.partition_image(
|
||||
example_doc_path("img/layout-parser-paper-with-table.jpg"),
|
||||
strategy=strategy,
|
||||
starting_page_number=2,
|
||||
)
|
||||
ids_1 = [element.id for element in elements_1]
|
||||
ids_2 = [element.id for element in elements_2]
|
||||
|
||||
assert ids_1 == ids_2
|
||||
|
||||
|
||||
def test_multi_page_tiff_starts_on_starting_page_number():
|
||||
elements = image.partition_image(
|
||||
example_doc_path("img/layout-parser-paper-combined.tiff"),
|
||||
starting_page_number=2,
|
||||
)
|
||||
pages = {element.metadata.page_number for element in elements}
|
||||
|
||||
assert pages == {2, 3}
|
||||
@@ -0,0 +1,169 @@
|
||||
from unstructured_inference.inference.elements import TextRegion, TextRegions
|
||||
from unstructured_inference.inference.layoutelement import LayoutElement, LayoutElements
|
||||
|
||||
from unstructured.documents.elements import ElementType
|
||||
from unstructured.partition.pdf_image.inference_utils import (
|
||||
build_layout_elements_from_ocr_regions,
|
||||
merge_text_regions,
|
||||
)
|
||||
|
||||
|
||||
def test_merge_text_regions(mock_embedded_text_regions):
|
||||
expected = TextRegion.from_coords(
|
||||
x1=437.83888888888885,
|
||||
y1=317.319341111111,
|
||||
x2=1256.334784222222,
|
||||
y2=406.9837855555556,
|
||||
text="LayoutParser: A Unified Toolkit for Deep Learning Based Document Image",
|
||||
)
|
||||
|
||||
merged_text_region = merge_text_regions(TextRegions.from_list(mock_embedded_text_regions))
|
||||
assert merged_text_region == expected
|
||||
|
||||
|
||||
def test_build_layout_elements_from_ocr_regions(mock_embedded_text_regions):
|
||||
expected = LayoutElements.from_list(
|
||||
[
|
||||
LayoutElement.from_coords(
|
||||
x1=437.83888888888885,
|
||||
y1=317.319341111111,
|
||||
x2=1256.334784222222,
|
||||
y2=406.9837855555556,
|
||||
text="LayoutParser: A Unified Toolkit for Deep Learning Based Document Image",
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
elements = build_layout_elements_from_ocr_regions(
|
||||
TextRegions.from_list(mock_embedded_text_regions)
|
||||
)
|
||||
assert elements == expected
|
||||
|
||||
|
||||
def test_build_layout_elements_from_ocr_regions_with_text(mock_embedded_text_regions):
|
||||
text = "LayoutParser: A Unified Toolkit for Deep Learning Based Document Image"
|
||||
expected = LayoutElements.from_list(
|
||||
[
|
||||
LayoutElement.from_coords(
|
||||
x1=437.83888888888885,
|
||||
y1=317.319341111111,
|
||||
x2=1256.334784222222,
|
||||
y2=406.9837855555556,
|
||||
text=text,
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
elements = build_layout_elements_from_ocr_regions(
|
||||
TextRegions.from_list(mock_embedded_text_regions),
|
||||
text,
|
||||
group_by_ocr_text=True,
|
||||
)
|
||||
assert elements == expected
|
||||
|
||||
|
||||
def test_build_layout_elements_from_ocr_regions_with_multi_line_text(mock_embedded_text_regions):
|
||||
text = "LayoutParser: \n\nA Unified Toolkit for Deep Learning Based Document Image"
|
||||
elements = build_layout_elements_from_ocr_regions(
|
||||
TextRegions.from_list(mock_embedded_text_regions),
|
||||
text,
|
||||
group_by_ocr_text=True,
|
||||
)
|
||||
assert elements == LayoutElements.from_list(
|
||||
[
|
||||
LayoutElement.from_coords(
|
||||
x1=453.00277777777774,
|
||||
y1=317.319341111111,
|
||||
x2=711.5338541666665,
|
||||
y2=358.28571222222206,
|
||||
text="LayoutParser:",
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
),
|
||||
LayoutElement.from_coords(
|
||||
x1=437.83888888888885,
|
||||
y1=317.319341111111,
|
||||
x2=1256.334784222222,
|
||||
y2=406.9837855555556,
|
||||
text="A Unified Toolkit for Deep Learning Based Document Image",
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_build_layout_elements_from_ocr_regions_with_repeated_texts(mock_embedded_text_regions):
|
||||
mock_embedded_text_regions.extend(
|
||||
[
|
||||
LayoutElement.from_coords(
|
||||
x1=453.00277777777774,
|
||||
y1=417.319341111111,
|
||||
x2=711.5338541666665,
|
||||
y2=458.28571222222206,
|
||||
text="LayoutParser",
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
),
|
||||
LayoutElement.from_coords(
|
||||
x1=453.00277777777774,
|
||||
y1=468.319341111111,
|
||||
x2=711.5338541666665,
|
||||
y2=478.28571222222206,
|
||||
text="for",
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
),
|
||||
LayoutElement.from_coords(
|
||||
x1=453.00277777777774,
|
||||
y1=488.319341111111,
|
||||
x2=711.5338541666665,
|
||||
y2=500.28571222222206,
|
||||
text="Deep",
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
),
|
||||
LayoutElement.from_coords(
|
||||
x1=453.00277777777774,
|
||||
y1=510.319341111111,
|
||||
x2=711.5338541666665,
|
||||
y2=550.28571222222206,
|
||||
text="Learning",
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
),
|
||||
]
|
||||
)
|
||||
text = (
|
||||
"LayoutParser: \n\nA Unified Toolkit for Deep Learning Based Document Image\n\n"
|
||||
"LayoutParser for Deep Learning"
|
||||
)
|
||||
elements = build_layout_elements_from_ocr_regions(
|
||||
TextRegions.from_list(mock_embedded_text_regions),
|
||||
text,
|
||||
group_by_ocr_text=True,
|
||||
)
|
||||
assert elements == LayoutElements.from_list(
|
||||
[
|
||||
LayoutElement.from_coords(
|
||||
x1=453.00277777777774,
|
||||
y1=317.319341111111,
|
||||
x2=711.5338541666665,
|
||||
y2=358.28571222222206,
|
||||
text="LayoutParser:",
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
),
|
||||
LayoutElement.from_coords(
|
||||
x1=437.83888888888885,
|
||||
y1=317.319341111111,
|
||||
x2=1256.334784222222,
|
||||
y2=406.9837855555556,
|
||||
text="A Unified Toolkit for Deep Learning Based Document Image",
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
),
|
||||
LayoutElement.from_coords(
|
||||
x1=453.00277777777774,
|
||||
y1=417.319341111111,
|
||||
x2=711.5338541666665,
|
||||
y2=550.28571222222206,
|
||||
text="LayoutParser for Deep Learning",
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
),
|
||||
]
|
||||
)
|
||||
@@ -0,0 +1,48 @@
|
||||
from PIL import Image
|
||||
from unstructured_inference.constants import IsExtracted
|
||||
from unstructured_inference.inference.elements import Rectangle
|
||||
from unstructured_inference.inference.layout import DocumentLayout, PageLayout
|
||||
from unstructured_inference.inference.layoutelement import LayoutElement, LayoutElements
|
||||
|
||||
from unstructured.partition.pdf_image.pdfminer_processing import (
|
||||
merge_inferred_with_extracted_layout,
|
||||
)
|
||||
|
||||
|
||||
def test_text_source_preserved_during_merge():
|
||||
"""Test that text_source property is preserved when elements are merged."""
|
||||
|
||||
# Create two simple LayoutElements with different text_source values
|
||||
inferred_element = LayoutElement(
|
||||
bbox=Rectangle(0, 0, 100, 50), text=None, is_extracted=IsExtracted.FALSE
|
||||
)
|
||||
|
||||
extracted_element = LayoutElement(
|
||||
bbox=Rectangle(0, 0, 100, 50), text="Extracted text", is_extracted=IsExtracted.TRUE
|
||||
)
|
||||
|
||||
# Create LayoutElements arrays
|
||||
inferred_layout_elements = LayoutElements.from_list([inferred_element])
|
||||
extracted_layout_elements = LayoutElements.from_list([extracted_element])
|
||||
|
||||
# Create a PageLayout for the inferred layout
|
||||
image = Image.new("RGB", (200, 200))
|
||||
inferred_page = PageLayout(number=1, image=image)
|
||||
inferred_page.elements_array = inferred_layout_elements
|
||||
|
||||
# Create DocumentLayout from the PageLayout
|
||||
inferred_document_layout = DocumentLayout(pages=[inferred_page])
|
||||
|
||||
# Merge them
|
||||
merged_layout = merge_inferred_with_extracted_layout(
|
||||
inferred_document_layout=inferred_document_layout,
|
||||
extracted_layout=[extracted_layout_elements],
|
||||
hi_res_model_name="test_model",
|
||||
)
|
||||
|
||||
# Verify text_source is preserved
|
||||
# Check the merged page's elements_array
|
||||
merged_page = merged_layout.pages[0]
|
||||
assert "Extracted text" in merged_page.elements_array.texts
|
||||
assert hasattr(merged_page.elements_array, "is_extracted_array")
|
||||
assert IsExtracted.TRUE in merged_page.elements_array.is_extracted_array
|
||||
@@ -0,0 +1,675 @@
|
||||
from collections import namedtuple
|
||||
from typing import Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
import unstructured_pytesseract
|
||||
from lxml import etree
|
||||
from pdf2image.exceptions import PDFPageCountError
|
||||
from PIL import Image, UnidentifiedImageError
|
||||
from unstructured_inference.inference.elements import EmbeddedTextRegion, TextRegion, TextRegions
|
||||
from unstructured_inference.inference.layout import DocumentLayout, PageLayout
|
||||
from unstructured_inference.inference.layoutelement import (
|
||||
LayoutElement,
|
||||
LayoutElements,
|
||||
)
|
||||
|
||||
from test_unstructured.unit_utils import example_doc_path
|
||||
from unstructured.documents.elements import ElementType
|
||||
from unstructured.partition.pdf_image import ocr
|
||||
from unstructured.partition.pdf_image.pdf_image_utils import (
|
||||
convert_pdf_to_images,
|
||||
pad_element_bboxes,
|
||||
)
|
||||
from unstructured.partition.utils.config import env_config
|
||||
from unstructured.partition.utils.constants import (
|
||||
OCR_AGENT_PADDLE,
|
||||
OCR_AGENT_TESSERACT,
|
||||
Source,
|
||||
)
|
||||
from unstructured.partition.utils.ocr_models.google_vision_ocr import OCRAgentGoogleVision
|
||||
from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
|
||||
from unstructured.partition.utils.ocr_models.paddle_ocr import OCRAgentPaddle
|
||||
from unstructured.partition.utils.ocr_models.tesseract_ocr import (
|
||||
OCRAgentTesseract,
|
||||
zoom_image,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("is_image", "expected_error"),
|
||||
[
|
||||
(True, UnidentifiedImageError),
|
||||
(False, PDFPageCountError),
|
||||
],
|
||||
)
|
||||
def test_process_data_with_ocr_invalid_file(is_image, expected_error):
|
||||
invalid_data = b"i am not a valid file"
|
||||
with pytest.raises(expected_error):
|
||||
_ = ocr.process_data_with_ocr(
|
||||
data=invalid_data,
|
||||
is_image=is_image,
|
||||
out_layout=DocumentLayout(),
|
||||
extracted_layout=[],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_image", [True, False])
|
||||
def test_process_file_with_ocr_invalid_filename(is_image):
|
||||
invalid_filename = "i am not a valid file name"
|
||||
with pytest.raises(FileNotFoundError):
|
||||
_ = ocr.process_file_with_ocr(
|
||||
filename=invalid_filename,
|
||||
is_image=is_image,
|
||||
out_layout=DocumentLayout(),
|
||||
extracted_layout=[],
|
||||
)
|
||||
|
||||
|
||||
def test_supplement_page_layout_with_ocr_invalid_ocr():
|
||||
with pytest.raises(ValueError):
|
||||
_ = ocr.supplement_page_layout_with_ocr(
|
||||
page_layout=None, image=None, ocr_agent="invliad_ocr"
|
||||
)
|
||||
|
||||
|
||||
def test_get_ocr_layout_from_image_tesseract(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
OCRAgentTesseract,
|
||||
"image_to_data_with_character_confidence_filter",
|
||||
lambda *args, **kwargs: pd.DataFrame(
|
||||
{
|
||||
"left": [10, 20, 30, 0],
|
||||
"top": [5, 15, 25, 0],
|
||||
"width": [15, 25, 35, 0],
|
||||
"height": [10, 20, 30, 0],
|
||||
"text": ["Hello", "World", "!", ""],
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
image = Image.new("RGB", (100, 100))
|
||||
|
||||
ocr_agent = OCRAgentTesseract()
|
||||
ocr_layout = ocr_agent.get_layout_from_image(image)
|
||||
|
||||
expected_layout = TextRegions(
|
||||
element_coords=np.array([[10.0, 5, 25, 15], [20, 15, 45, 35], [30, 25, 65, 55]]),
|
||||
texts=np.array(["Hello", "World", "!"]),
|
||||
sources=np.array([Source.OCR_TESSERACT] * 3),
|
||||
)
|
||||
|
||||
assert ocr_layout.texts.tolist() == expected_layout.texts.tolist()
|
||||
np.testing.assert_array_equal(ocr_layout.element_coords, expected_layout.element_coords)
|
||||
np.testing.assert_array_equal(ocr_layout.sources, expected_layout.sources)
|
||||
|
||||
|
||||
def mock_ocr(*args, **kwargs):
|
||||
return [
|
||||
[
|
||||
(
|
||||
[(10, 5), (25, 5), (25, 15), (10, 15)],
|
||||
["Hello"],
|
||||
),
|
||||
],
|
||||
[
|
||||
(
|
||||
[(20, 15), (45, 15), (45, 35), (20, 35)],
|
||||
["World"],
|
||||
),
|
||||
],
|
||||
[
|
||||
(
|
||||
[(30, 25), (65, 25), (65, 55), (30, 55)],
|
||||
["!"],
|
||||
),
|
||||
],
|
||||
[
|
||||
(
|
||||
[(0, 0), (0, 0), (0, 0), (0, 0)],
|
||||
[""],
|
||||
),
|
||||
],
|
||||
]
|
||||
|
||||
|
||||
def monkeypatch_load_agent(*args):
|
||||
class MockAgent:
|
||||
def __init__(self):
|
||||
self.ocr = mock_ocr
|
||||
|
||||
return MockAgent()
|
||||
|
||||
|
||||
def test_get_ocr_layout_from_image_paddle(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
OCRAgentPaddle,
|
||||
"load_agent",
|
||||
monkeypatch_load_agent,
|
||||
)
|
||||
|
||||
image = Image.new("RGB", (100, 100))
|
||||
|
||||
ocr_layout = OCRAgentPaddle().get_layout_from_image(image)
|
||||
|
||||
expected_layout = TextRegions(
|
||||
element_coords=np.array([[10.0, 5, 25, 15], [20, 15, 45, 35], [30, 25, 65, 55]]),
|
||||
texts=np.array(["Hello", "World", "!"]),
|
||||
sources=np.array([Source.OCR_PADDLE] * 3),
|
||||
)
|
||||
|
||||
assert ocr_layout.texts.tolist() == expected_layout.texts.tolist()
|
||||
np.testing.assert_array_equal(ocr_layout.element_coords, expected_layout.element_coords)
|
||||
np.testing.assert_array_equal(ocr_layout.sources, expected_layout.sources)
|
||||
|
||||
|
||||
def test_get_ocr_text_from_image_tesseract(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
unstructured_pytesseract,
|
||||
"image_to_string",
|
||||
lambda *args, **kwargs: "Hello World",
|
||||
)
|
||||
image = Image.new("RGB", (100, 100))
|
||||
|
||||
ocr_agent = OCRAgentTesseract()
|
||||
ocr_text = ocr_agent.get_text_from_image(image)
|
||||
|
||||
assert ocr_text == "Hello World"
|
||||
|
||||
|
||||
def test_get_ocr_text_from_image_paddle(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
OCRAgentPaddle,
|
||||
"load_agent",
|
||||
monkeypatch_load_agent,
|
||||
)
|
||||
|
||||
image = Image.new("RGB", (100, 100))
|
||||
|
||||
ocr_agent = OCRAgentPaddle()
|
||||
ocr_text = ocr_agent.get_text_from_image(image)
|
||||
|
||||
assert ocr_text == "Hello\n\nWorld\n\n!"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def google_vision_text_annotation():
|
||||
from google.cloud.vision import (
|
||||
Block,
|
||||
BoundingPoly,
|
||||
Page,
|
||||
Paragraph,
|
||||
Symbol,
|
||||
TextAnnotation,
|
||||
Vertex,
|
||||
Word,
|
||||
)
|
||||
|
||||
breaks = TextAnnotation.DetectedBreak.BreakType
|
||||
symbols_hello = [Symbol(text=c) for c in "Hello"] + [
|
||||
Symbol(
|
||||
property=TextAnnotation.TextProperty(
|
||||
detected_break=TextAnnotation.DetectedBreak(type_=breaks.SPACE)
|
||||
)
|
||||
)
|
||||
]
|
||||
symbols_world = [Symbol(text=c) for c in "World!"] + [
|
||||
Symbol(
|
||||
property=TextAnnotation.TextProperty(
|
||||
detected_break=TextAnnotation.DetectedBreak(type_=breaks.LINE_BREAK)
|
||||
)
|
||||
)
|
||||
]
|
||||
words = [Word(symbols=symbols_hello), Word(symbols=symbols_world)]
|
||||
bounding_box = BoundingPoly(
|
||||
vertices=[Vertex(x=0, y=0), Vertex(x=0, y=10), Vertex(x=10, y=10), Vertex(x=10, y=0)]
|
||||
)
|
||||
paragraphs = [Paragraph(words=words, bounding_box=bounding_box)]
|
||||
blocks = [Block(paragraphs=paragraphs)]
|
||||
pages = [Page(blocks=blocks)]
|
||||
return TextAnnotation(text="Hello World!", pages=pages)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def google_vision_client(google_vision_text_annotation):
|
||||
Response = namedtuple("Response", "full_text_annotation")
|
||||
|
||||
class FakeGoogleVisionClient:
|
||||
def document_text_detection(self, image, image_context):
|
||||
return Response(full_text_annotation=google_vision_text_annotation)
|
||||
|
||||
class OCRAgentFakeGoogleVision(OCRAgentGoogleVision):
|
||||
def __init__(self, language: Optional[str] = None):
|
||||
self.client = FakeGoogleVisionClient()
|
||||
self.language = language
|
||||
|
||||
return OCRAgentFakeGoogleVision()
|
||||
|
||||
|
||||
def test_get_ocr_from_image_google_vision(google_vision_client):
|
||||
image = Image.new("RGB", (100, 100))
|
||||
|
||||
ocr_agent = google_vision_client
|
||||
ocr_text = ocr_agent.get_text_from_image(image)
|
||||
|
||||
assert ocr_text == "Hello World!"
|
||||
|
||||
|
||||
def test_get_layout_from_image_google_vision(google_vision_client):
|
||||
image = Image.new("RGB", (100, 100))
|
||||
|
||||
ocr_agent = google_vision_client
|
||||
regions = ocr_agent.get_layout_from_image(image)
|
||||
assert len(regions) == 1
|
||||
assert regions.texts[0] == "Hello World!"
|
||||
assert all(source == Source.OCR_GOOGLEVISION for source in regions.sources)
|
||||
assert regions.x1[0] == 0
|
||||
assert regions.y1[0] == 0
|
||||
assert regions.x2[0] == 10
|
||||
assert regions.y2[0] == 10
|
||||
|
||||
|
||||
def test_get_layout_elements_from_image_google_vision(google_vision_client):
|
||||
image = Image.new("RGB", (100, 100))
|
||||
|
||||
ocr_agent = google_vision_client
|
||||
layout_elements = ocr_agent.get_layout_elements_from_image(image)
|
||||
assert len(layout_elements) == 1
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_ocr_regions():
|
||||
return TextRegions.from_list(
|
||||
[
|
||||
EmbeddedTextRegion.from_coords(10, 10, 90, 90, text="0", source=None),
|
||||
EmbeddedTextRegion.from_coords(200, 200, 300, 300, text="1", source=None),
|
||||
EmbeddedTextRegion.from_coords(500, 320, 600, 350, text="3", source=None),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_out_layout(mock_embedded_text_regions):
|
||||
return LayoutElements.from_list(
|
||||
[
|
||||
LayoutElement(
|
||||
text="",
|
||||
source=None,
|
||||
type="Text",
|
||||
bbox=r.bbox,
|
||||
)
|
||||
for r in mock_embedded_text_regions
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_aggregate_ocr_text_by_block():
|
||||
expected = "A Unified Toolkit"
|
||||
ocr_layout = [
|
||||
TextRegion.from_coords(0, 0, 20, 20, "A"),
|
||||
TextRegion.from_coords(50, 50, 150, 150, "Unified"),
|
||||
TextRegion.from_coords(150, 150, 300, 250, "Toolkit"),
|
||||
TextRegion.from_coords(200, 250, 300, 350, "Deep"),
|
||||
]
|
||||
region = TextRegion.from_coords(0, 0, 250, 350, "")
|
||||
|
||||
text = ocr.aggregate_ocr_text_by_block(ocr_layout, region, 0.5)
|
||||
assert text == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("zoom", [1, 0.1, 5, -1, 0])
|
||||
def test_zoom_image(zoom):
|
||||
image = Image.new("RGB", (100, 100))
|
||||
width, height = image.size
|
||||
new_image = zoom_image(image, zoom)
|
||||
new_w, new_h = new_image.size
|
||||
if zoom <= 0:
|
||||
zoom = 1
|
||||
assert new_w == np.round(width * zoom, 0)
|
||||
assert new_h == np.round(height * zoom, 0)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_layout(mock_embedded_text_regions):
|
||||
return LayoutElements.from_list(
|
||||
[
|
||||
LayoutElement(text=r.text, type=ElementType.UNCATEGORIZED_TEXT, bbox=r.bbox)
|
||||
for r in mock_embedded_text_regions
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_supplement_layout_with_ocr_elements(mock_layout, mock_ocr_regions):
|
||||
ocr_elements = [
|
||||
LayoutElement(text=r.text, source=None, type=ElementType.UNCATEGORIZED_TEXT, bbox=r.bbox)
|
||||
for r in mock_ocr_regions.as_list()
|
||||
]
|
||||
|
||||
final_layout = ocr.supplement_layout_with_ocr_elements(mock_layout, mock_ocr_regions).as_list()
|
||||
|
||||
# Check if the final layout contains the original layout elements
|
||||
for element in mock_layout.as_list():
|
||||
assert element in final_layout
|
||||
|
||||
# Check if the final layout contains the OCR-derived elements
|
||||
assert any(ocr_element in final_layout for ocr_element in ocr_elements)
|
||||
|
||||
# Check if the OCR-derived elements that are subregions of layout elements are removed
|
||||
for element in mock_layout.as_list():
|
||||
for ocr_element in ocr_elements:
|
||||
if ocr_element.bbox.is_almost_subregion_of(
|
||||
element.bbox,
|
||||
env_config.OCR_LAYOUT_SUBREGION_THRESHOLD,
|
||||
):
|
||||
assert ocr_element not in final_layout
|
||||
|
||||
|
||||
def test_merge_out_layout_with_ocr_layout(mock_out_layout, mock_ocr_regions):
|
||||
ocr_elements = [
|
||||
LayoutElement(text=r.text, source=None, type=ElementType.UNCATEGORIZED_TEXT, bbox=r.bbox)
|
||||
for r in mock_ocr_regions.as_list()
|
||||
]
|
||||
input_layout_elements = mock_out_layout.as_list()
|
||||
|
||||
final_layout = ocr.merge_out_layout_with_ocr_layout(
|
||||
mock_out_layout,
|
||||
mock_ocr_regions,
|
||||
).as_list()
|
||||
|
||||
# Check if the out layout's text attribute is updated with aggregated OCR text
|
||||
assert final_layout[0].text == mock_ocr_regions.texts[2]
|
||||
|
||||
# Check if the final layout contains both original elements and OCR-derived elements
|
||||
# The first element's text is modified by the ocr regions so it won't be the same as the input
|
||||
assert all(element in final_layout for element in input_layout_elements[1:])
|
||||
assert final_layout[0].bbox == input_layout_elements[0].bbox
|
||||
assert any(element in final_layout for element in ocr_elements)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("padding", "expected_bbox"),
|
||||
[
|
||||
(5, (5, 15, 35, 45)),
|
||||
(-3, (13, 23, 27, 37)),
|
||||
(2.5, (7.5, 17.5, 32.5, 42.5)),
|
||||
(-1.5, (11.5, 21.5, 28.5, 38.5)),
|
||||
],
|
||||
)
|
||||
def test_pad_element_bboxes(padding, expected_bbox):
|
||||
element = LayoutElement.from_coords(
|
||||
x1=10,
|
||||
y1=20,
|
||||
x2=30,
|
||||
y2=40,
|
||||
text="",
|
||||
source=None,
|
||||
type=ElementType.UNCATEGORIZED_TEXT,
|
||||
)
|
||||
expected_original_element_bbox = (10, 20, 30, 40)
|
||||
|
||||
padded_element = pad_element_bboxes(element, padding)
|
||||
|
||||
padded_element_bbox = (
|
||||
padded_element.bbox.x1,
|
||||
padded_element.bbox.y1,
|
||||
padded_element.bbox.x2,
|
||||
padded_element.bbox.y2,
|
||||
)
|
||||
assert padded_element_bbox == expected_bbox
|
||||
|
||||
# make sure the original element has not changed
|
||||
original_element_bbox = (element.bbox.x1, element.bbox.y1, element.bbox.x2, element.bbox.y2)
|
||||
assert original_element_bbox == expected_original_element_bbox
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def table_element():
|
||||
table = LayoutElement.from_coords(x1=10, y1=20, x2=50, y2=70, text="I am a table", type="Table")
|
||||
return table
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_ocr_layout():
|
||||
return TextRegions.from_list(
|
||||
[
|
||||
TextRegion.from_coords(x1=15, y1=25, x2=35, y2=45, text="Token1"),
|
||||
TextRegion.from_coords(x1=40, y1=30, x2=45, y2=50, text="Token2"),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_supplement_element_with_table_extraction():
|
||||
from unstructured_inference.models import tables
|
||||
|
||||
tables.load_agent()
|
||||
|
||||
image = next(convert_pdf_to_images(example_doc_path("pdf/single_table.pdf")))
|
||||
elements = LayoutElements(
|
||||
element_coords=np.array([[215.00109863, 731.89996338, 1470.07739258, 972.83129883]]),
|
||||
texts=np.array(["foo"]),
|
||||
sources=np.array(["yolox_sg"]),
|
||||
element_class_ids=np.array([0]),
|
||||
element_class_id_map={0: "Table"},
|
||||
)
|
||||
supplemented = ocr.supplement_element_with_table_extraction(
|
||||
elements=elements,
|
||||
image=image,
|
||||
tables_agent=tables.tables_agent,
|
||||
ocr_agent=ocr.OCRAgent.get_agent(language="eng"),
|
||||
)
|
||||
assert supplemented.text_as_html[0].startswith("<table>")
|
||||
|
||||
|
||||
def test_get_table_tokens(mock_ocr_layout):
|
||||
with patch.object(OCRAgentTesseract, "get_layout_from_image", return_value=mock_ocr_layout):
|
||||
ocr_agent = OCRAgent.get_agent(language="eng")
|
||||
table_tokens = ocr.get_table_tokens(table_element_image=None, ocr_agent=ocr_agent)
|
||||
expected_tokens = [
|
||||
{
|
||||
"bbox": [15, 25, 35, 45],
|
||||
"text": "Token1",
|
||||
"span_num": 0,
|
||||
"line_num": 0,
|
||||
"block_num": 0,
|
||||
},
|
||||
{
|
||||
"bbox": [40, 30, 45, 50],
|
||||
"text": "Token2",
|
||||
"span_num": 1,
|
||||
"line_num": 0,
|
||||
"block_num": 0,
|
||||
},
|
||||
]
|
||||
|
||||
assert table_tokens == expected_tokens
|
||||
|
||||
|
||||
def test_auto_zoom_not_exceed_tesseract_limit(monkeypatch):
|
||||
monkeypatch.setenv("TESSERACT_MIN_TEXT_HEIGHT", "1000")
|
||||
monkeypatch.setenv("TESSERACT_OPTIMUM_TEXT_HEIGHT", "100000")
|
||||
monkeypatch.setattr(
|
||||
OCRAgentTesseract,
|
||||
"image_to_data_with_character_confidence_filter",
|
||||
lambda *args, **kwargs: pd.DataFrame(
|
||||
{
|
||||
"left": [10, 20, 30, 0],
|
||||
"top": [5, 15, 25, 0],
|
||||
"width": [15, 25, 35, 0],
|
||||
"height": [10, 20, 30, 0],
|
||||
"text": ["Hello", "World", "!", ""],
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
image = Image.new("RGB", (1000, 1000))
|
||||
ocr_agent = OCRAgentTesseract()
|
||||
# tests that the code can run instead of oom and OCR results make sense
|
||||
assert ocr_agent.get_layout_from_image(image).texts.tolist() == [
|
||||
"Hello",
|
||||
"World",
|
||||
"!",
|
||||
]
|
||||
|
||||
|
||||
def test_merge_out_layout_with_cid_code(mock_out_layout, mock_ocr_regions):
|
||||
# the code should ignore this invalid text and use ocr region's text
|
||||
mock_out_layout.texts = mock_out_layout.texts.astype(object)
|
||||
mock_out_layout.texts[0] = "(cid:10)(cid:5)?"
|
||||
ocr_elements = [
|
||||
LayoutElement(text=r.text, source=None, type=ElementType.UNCATEGORIZED_TEXT, bbox=r.bbox)
|
||||
for r in mock_ocr_regions.as_list()
|
||||
]
|
||||
input_layout_elements = mock_out_layout.as_list()
|
||||
|
||||
# TODO (yao): refactor the tests to check the array data structure directly instead of
|
||||
# converting them into lists first (this includes other tests in this file)
|
||||
final_layout = ocr.merge_out_layout_with_ocr_layout(mock_out_layout, mock_ocr_regions).as_list()
|
||||
|
||||
# Check if the out layout's text attribute is updated with aggregated OCR text
|
||||
assert final_layout[0].text == mock_ocr_regions.texts[2]
|
||||
|
||||
# Check if the final layout contains both original elements and OCR-derived elements
|
||||
assert all(element in final_layout for element in input_layout_elements[1:])
|
||||
assert any(element in final_layout for element in ocr_elements)
|
||||
|
||||
|
||||
def _create_hocr_word_span(
|
||||
characters: list[tuple[str, str]], word_bbox: tuple[int, int, int, int], namespace_map: dict
|
||||
) -> etree.Element:
|
||||
word_span = [
|
||||
'<root xmlns="http://www.w3.org/1999/xhtml" xml:lang="en" lang="en">\n',
|
||||
(
|
||||
f"<span class='ocrx_word' title='"
|
||||
f"bbox {word_bbox[0]} {word_bbox[1]} {word_bbox[2]} {word_bbox[3]}"
|
||||
f"; x_wconf 64'>"
|
||||
),
|
||||
]
|
||||
for char, x_conf in characters:
|
||||
word_span.append(
|
||||
f"<span class='ocrx_cinfo' title='x_bboxes 0 0 0 0; x_conf {x_conf}'>{char}</span>"
|
||||
)
|
||||
word_span.append("</span>")
|
||||
word_span.append("</root>")
|
||||
root = etree.fromstring("\n".join(word_span))
|
||||
return root
|
||||
|
||||
|
||||
def test_extract_word_from_hocr():
|
||||
characters = [
|
||||
("w", "99.0"),
|
||||
("o", "98.5"),
|
||||
("r", "97.5"),
|
||||
("d", "96.0"),
|
||||
("!", "50.0"),
|
||||
("@", "45.0"),
|
||||
]
|
||||
word_bbox = (10, 9, 70, 22)
|
||||
agent = OCRAgentTesseract()
|
||||
word_span = _create_hocr_word_span(characters, word_bbox, agent.hocr_namespace)
|
||||
|
||||
text = agent.extract_word_from_hocr(word_span, 0.0)
|
||||
assert text == "word!@"
|
||||
|
||||
text = agent.extract_word_from_hocr(word_span, 0.960)
|
||||
assert text == "word"
|
||||
|
||||
text = agent.extract_word_from_hocr(word_span, 0.990)
|
||||
assert text == "w"
|
||||
|
||||
text = agent.extract_word_from_hocr(word_span, 0.999)
|
||||
assert text == ""
|
||||
|
||||
|
||||
def test_hocr_to_dataframe():
|
||||
characters = [
|
||||
("w", "99.0"),
|
||||
("o", "98.5"),
|
||||
("r", "97.5"),
|
||||
("d", "96.0"),
|
||||
("!", "50.0"),
|
||||
("@", "45.0"),
|
||||
]
|
||||
word_bbox = (10, 9, 70, 22)
|
||||
agent = OCRAgentTesseract()
|
||||
hocr = etree.tostring(_create_hocr_word_span(characters, word_bbox, agent.hocr_namespace))
|
||||
df = agent.hocr_to_dataframe(hocr=hocr, character_confidence_threshold=0.960)
|
||||
|
||||
assert df.shape == (1, 5)
|
||||
assert df["left"].iloc[0] == 10
|
||||
assert df["top"].iloc[0] == 9
|
||||
assert df["width"].iloc[0] == 60
|
||||
assert df["height"].iloc[0] == 13
|
||||
assert df["text"].iloc[0] == "word"
|
||||
|
||||
|
||||
def test_hocr_to_dataframe_when_no_prediction_empty_df():
|
||||
df = OCRAgentTesseract().hocr_to_dataframe(hocr="")
|
||||
|
||||
assert df.shape == (0, 5)
|
||||
assert "left" in df.columns
|
||||
assert "top" in df.columns
|
||||
assert "width" in df.columns
|
||||
assert "height" in df.columns
|
||||
assert "text" in df.columns
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_page(mock_ocr_layout, mock_layout):
|
||||
mock_page = MagicMock(PageLayout)
|
||||
mock_page.elements_array = mock_layout
|
||||
return mock_page
|
||||
|
||||
|
||||
def test_supplement_layout_with_ocr(mock_ocr_get_instance, mocker, mock_page):
|
||||
from unstructured.partition.pdf_image.ocr import OCRAgent
|
||||
|
||||
mocker.patch.object(OCRAgent, "get_layout_from_image", return_value=mock_ocr_layout)
|
||||
|
||||
ocr.supplement_page_layout_with_ocr(
|
||||
mock_page,
|
||||
Image.new("RGB", (100, 100)),
|
||||
infer_table_structure=True,
|
||||
ocr_agent=OCR_AGENT_TESSERACT,
|
||||
ocr_languages="eng",
|
||||
table_ocr_agent=OCR_AGENT_PADDLE,
|
||||
)
|
||||
|
||||
assert mock_ocr_get_instance.call_args_list[0][1] == {
|
||||
"language": "eng",
|
||||
"ocr_agent_module": OCR_AGENT_TESSERACT,
|
||||
}
|
||||
assert mock_ocr_get_instance.call_args_list[1][1] == {
|
||||
"language": "en",
|
||||
"ocr_agent_module": OCR_AGENT_PADDLE,
|
||||
}
|
||||
|
||||
|
||||
def test_pass_down_agents(mock_ocr_get_instance, mocker, mock_page):
|
||||
from unstructured.partition.pdf_image.ocr import OCRAgent, PILImage
|
||||
|
||||
mocker.patch.object(OCRAgent, "get_layout_from_image", return_value=mock_ocr_layout)
|
||||
mocker.patch.object(PILImage, "open", return_value=Image.new("RGB", (100, 100)))
|
||||
doc = MagicMock(DocumentLayout)
|
||||
doc.pages = [mock_page]
|
||||
|
||||
ocr.process_file_with_ocr(
|
||||
"foo",
|
||||
doc,
|
||||
[],
|
||||
infer_table_structure=True,
|
||||
is_image=True,
|
||||
ocr_agent=OCR_AGENT_PADDLE,
|
||||
ocr_languages="eng",
|
||||
table_ocr_agent=OCR_AGENT_TESSERACT,
|
||||
)
|
||||
|
||||
assert mock_ocr_get_instance.call_args_list[0][1] == {
|
||||
"language": "en",
|
||||
"ocr_agent_module": OCR_AGENT_PADDLE,
|
||||
}
|
||||
assert mock_ocr_get_instance.call_args_list[1][1] == {
|
||||
"language": "eng",
|
||||
"ocr_agent_module": OCR_AGENT_TESSERACT,
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,383 @@
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image as PILImg
|
||||
|
||||
from test_unstructured.unit_utils import example_doc_path
|
||||
from unstructured.documents.coordinates import PixelSpace
|
||||
from unstructured.documents.elements import ElementMetadata, ElementType, Image, Table
|
||||
from unstructured.partition.pdf_image import pdf_image_utils
|
||||
|
||||
|
||||
@pytest.mark.parametrize("image_type", ["pil", "numpy_array"])
|
||||
def test_write_image(image_type):
|
||||
mock_pil_image = PILImg.new("RGB", (50, 50))
|
||||
mock_numpy_image = np.zeros((50, 50, 3), np.uint8)
|
||||
|
||||
image_map = {
|
||||
"pil": mock_pil_image,
|
||||
"numpy_array": mock_numpy_image,
|
||||
}
|
||||
image = image_map[image_type]
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
output_image_path = os.path.join(tmpdir, "test_image.jpg")
|
||||
pdf_image_utils.write_image(image, output_image_path)
|
||||
assert os.path.exists(output_image_path)
|
||||
|
||||
# Additional check to see if the written image can be read
|
||||
read_image = PILImg.open(output_image_path)
|
||||
assert read_image is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("file_mode", ["filename", "rb"])
|
||||
@pytest.mark.parametrize("path_only", [True, False])
|
||||
def test_convert_pdf_to_image(file_mode, path_only):
|
||||
filename = example_doc_path("pdf/embedded-images.pdf")
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
if file_mode == "filename":
|
||||
images = pdf_image_utils.convert_pdf_to_image(
|
||||
filename=filename,
|
||||
file=None,
|
||||
output_folder=tmpdir,
|
||||
path_only=path_only,
|
||||
)
|
||||
else:
|
||||
with open(filename, "rb") as f:
|
||||
images = pdf_image_utils.convert_pdf_to_image(
|
||||
filename="",
|
||||
file=f,
|
||||
output_folder=tmpdir,
|
||||
path_only=path_only,
|
||||
)
|
||||
|
||||
if path_only:
|
||||
assert isinstance(images[0], str)
|
||||
else:
|
||||
assert isinstance(images[0], PILImg.Image)
|
||||
|
||||
|
||||
def test_convert_pdf_to_image_raises_error():
|
||||
filename = example_doc_path("embedded-images.pdf")
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
pdf_image_utils.convert_pdf_to_image(filename=filename, path_only=True, output_folder=None)
|
||||
|
||||
assert str(exc_info.value) == "output_folder must be specified if path_only is true"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "is_image"),
|
||||
[
|
||||
(example_doc_path("pdf/layout-parser-paper-fast.pdf"), False),
|
||||
(example_doc_path("img/layout-parser-paper-fast.jpg"), True),
|
||||
(example_doc_path("img/english-and-korean.png"), True),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("element_category_to_save", [ElementType.IMAGE, ElementType.TABLE])
|
||||
@pytest.mark.parametrize("extract_image_block_to_payload", [False, True])
|
||||
@pytest.mark.parametrize("horizontal_padding", [0, 20])
|
||||
@pytest.mark.parametrize("vertical_padding", [0, 10])
|
||||
def test_save_elements(
|
||||
element_category_to_save,
|
||||
extract_image_block_to_payload,
|
||||
filename,
|
||||
is_image,
|
||||
horizontal_padding,
|
||||
vertical_padding,
|
||||
monkeypatch,
|
||||
):
|
||||
if horizontal_padding > 0:
|
||||
monkeypatch.setenv("EXTRACT_IMAGE_BLOCK_CROP_HORIZONTAL_PAD", str(horizontal_padding))
|
||||
if vertical_padding > 0:
|
||||
monkeypatch.setenv("EXTRACT_IMAGE_BLOCK_CROP_VERTICAL_PAD", str(vertical_padding))
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
elements = [
|
||||
Image(
|
||||
text="Image Text 1",
|
||||
coordinates=((78, 86), (78, 519), (512, 519), (512, 86)),
|
||||
coordinate_system=PixelSpace(width=1575, height=1166),
|
||||
metadata=ElementMetadata(page_number=1),
|
||||
),
|
||||
Image(
|
||||
text="Image Text 2",
|
||||
coordinates=((570, 86), (570, 519), (1003, 519), (1003, 86)),
|
||||
coordinate_system=PixelSpace(width=1575, height=1166),
|
||||
metadata=ElementMetadata(page_number=1),
|
||||
),
|
||||
Image(
|
||||
text="Table 1",
|
||||
coordinates=((1062, 86), (1062, 519), (1496, 519), (1496, 86)),
|
||||
coordinate_system=PixelSpace(width=1575, height=1166),
|
||||
metadata=ElementMetadata(page_number=1),
|
||||
),
|
||||
]
|
||||
if not is_image:
|
||||
# add a page 2 element
|
||||
elements.append(
|
||||
Table(
|
||||
text="Table 2",
|
||||
coordinates=((1062, 86), (1062, 519), (1496, 519), (1496, 86)),
|
||||
coordinate_system=PixelSpace(width=1575, height=1166),
|
||||
metadata=ElementMetadata(page_number=2),
|
||||
),
|
||||
)
|
||||
|
||||
pdf_image_utils.save_elements(
|
||||
elements=elements,
|
||||
starting_page_number=1,
|
||||
element_category_to_save=element_category_to_save,
|
||||
pdf_image_dpi=200,
|
||||
filename=filename,
|
||||
is_image=is_image,
|
||||
output_dir_path=str(tmpdir),
|
||||
extract_image_block_to_payload=extract_image_block_to_payload,
|
||||
)
|
||||
|
||||
saved_elements = [el for el in elements if el.category == element_category_to_save]
|
||||
for i, el in enumerate(saved_elements):
|
||||
basename = "table" if el.category == ElementType.TABLE else "figure"
|
||||
expected_image_path = os.path.join(
|
||||
str(tmpdir), f"{basename}-{el.metadata.page_number}-{i + 1}.jpg"
|
||||
)
|
||||
if extract_image_block_to_payload:
|
||||
assert isinstance(el.metadata.image_base64, str)
|
||||
assert isinstance(el.metadata.image_mime_type, str)
|
||||
image_bytes = base64.b64decode(el.metadata.image_base64)
|
||||
image = PILImg.open(io.BytesIO(image_bytes))
|
||||
x1, y1 = el.metadata.coordinates.points[0]
|
||||
x2, y2 = el.metadata.coordinates.points[2]
|
||||
width = x2 - x1
|
||||
height = y2 - y1
|
||||
assert image.width == width + 2 * horizontal_padding
|
||||
assert image.height == height + 2 * vertical_padding
|
||||
assert not el.metadata.image_path
|
||||
assert not os.path.isfile(expected_image_path)
|
||||
else:
|
||||
assert os.path.isfile(expected_image_path)
|
||||
image = PILImg.open(expected_image_path)
|
||||
x1, y1 = el.metadata.coordinates.points[0]
|
||||
x2, y2 = el.metadata.coordinates.points[2]
|
||||
width = x2 - x1
|
||||
height = y2 - y1
|
||||
assert image.width == width + 2 * horizontal_padding
|
||||
assert image.height == height + 2 * vertical_padding
|
||||
assert el.metadata.image_path == expected_image_path
|
||||
assert not el.metadata.image_base64
|
||||
assert not el.metadata.image_mime_type
|
||||
|
||||
|
||||
@pytest.mark.parametrize("storage_enabled", [False, True])
|
||||
def test_save_elements_with_output_dir_path_none(monkeypatch, storage_enabled):
|
||||
monkeypatch.setenv("GLOBAL_WORKING_DIR_ENABLED", storage_enabled)
|
||||
with (
|
||||
patch("PIL.Image.open"),
|
||||
patch("unstructured.partition.pdf_image.pdf_image_utils.write_image"),
|
||||
patch("unstructured.partition.pdf_image.pdf_image_utils.convert_pdf_to_image"),
|
||||
tempfile.TemporaryDirectory() as tmpdir,
|
||||
):
|
||||
original_cwd = os.getcwd()
|
||||
os.chdir(tmpdir)
|
||||
pdf_image_utils.save_elements(
|
||||
elements=[],
|
||||
element_category_to_save="",
|
||||
starting_page_number=1,
|
||||
pdf_image_dpi=200,
|
||||
filename="dummy.pdf",
|
||||
output_dir_path=None,
|
||||
)
|
||||
|
||||
# Verify that the images are saved in the expected directory
|
||||
if storage_enabled:
|
||||
from unstructured.partition.utils.config import env_config
|
||||
|
||||
expected_output_dir = os.path.join(env_config.GLOBAL_WORKING_PROCESS_DIR, "figures")
|
||||
else:
|
||||
expected_output_dir = os.path.join(tmpdir, "figures")
|
||||
assert os.path.exists(expected_output_dir)
|
||||
assert os.path.isdir(expected_output_dir)
|
||||
os.chdir(original_cwd)
|
||||
|
||||
|
||||
def test_write_image_raises_error():
|
||||
with pytest.raises(ValueError):
|
||||
pdf_image_utils.write_image("invalid_type", "test_image.jpg")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "outcome"), [("", False), ("foo", True), (None, False), ("(cid:10)boo", False)]
|
||||
)
|
||||
def test_valid_text(text, outcome):
|
||||
assert pdf_image_utils.valid_text(text) == outcome
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("base", 0.0),
|
||||
("", 0.0),
|
||||
("(cid:2)", 1.0),
|
||||
("(cid:1)a", 0.5),
|
||||
("c(cid:1)ab", 0.25),
|
||||
],
|
||||
)
|
||||
def test_cid_ratio(text, expected):
|
||||
assert pdf_image_utils.cid_ratio(text) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("base", False),
|
||||
("(cid:2)", True),
|
||||
("(cid:1234567890)", True),
|
||||
("jkl;(cid:12)asdf", True),
|
||||
],
|
||||
)
|
||||
def test_is_cid_present(text, expected):
|
||||
assert pdf_image_utils.is_cid_present(text) == expected
|
||||
|
||||
|
||||
def test_pad_bbox():
|
||||
bbox = (100, 100, 200, 200)
|
||||
padding = (10, 20) # Horizontal padding 10, Vertical padding 20
|
||||
expected = (90, 80, 210, 220)
|
||||
|
||||
result = pdf_image_utils.pad_bbox(bbox, padding)
|
||||
assert result == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("input_types", "expected"),
|
||||
[
|
||||
(None, []),
|
||||
(["table", "image"], ["Table", "Image"]),
|
||||
(["unknown"], ["Unknown"]),
|
||||
(["Table", "image", "UnknOwn"], ["Table", "Image", "Unknown"]),
|
||||
(["NarrativeText", "narrativetext"], ["NarrativeText", "NarrativeText"]),
|
||||
],
|
||||
)
|
||||
def test_check_element_types_to_extract(input_types, expected):
|
||||
assert pdf_image_utils.check_element_types_to_extract(input_types) == expected
|
||||
|
||||
|
||||
def test_check_element_types_to_extract_raises_error():
|
||||
with pytest.raises(TypeError) as exc_info:
|
||||
pdf_image_utils.check_element_types_to_extract("not a list")
|
||||
assert "must be a list" in str(exc_info.value)
|
||||
|
||||
|
||||
class MockPageLayout:
|
||||
def annotate(self, colors):
|
||||
return "mock_image"
|
||||
|
||||
|
||||
class MockDocumentLayout:
|
||||
pages = [MockPageLayout(), MockPageLayout]
|
||||
|
||||
|
||||
def test_annotate_layout_elements_with_image():
|
||||
inferred_layout = MockPageLayout()
|
||||
extracted_layout = MockPageLayout()
|
||||
output_basename = "test_page"
|
||||
page_number = 1
|
||||
|
||||
# Check if images for both layouts were saved
|
||||
with (
|
||||
tempfile.TemporaryDirectory() as tmpdir,
|
||||
patch("unstructured.partition.pdf_image.pdf_image_utils.write_image") as mock_write_image,
|
||||
):
|
||||
pdf_image_utils.annotate_layout_elements_with_image(
|
||||
inferred_page_layout=inferred_layout,
|
||||
extracted_page_layout=extracted_layout,
|
||||
output_dir_path=str(tmpdir),
|
||||
output_f_basename=output_basename,
|
||||
page_number=page_number,
|
||||
)
|
||||
|
||||
expected_filenames = [
|
||||
f"{output_basename}_{page_number}_inferred.jpg",
|
||||
f"{output_basename}_{page_number}_extracted.jpg",
|
||||
]
|
||||
actual_calls = [call.args[1] for call in mock_write_image.call_args_list]
|
||||
for expected_filename in expected_filenames:
|
||||
assert any(expected_filename in actual_call for actual_call in actual_calls)
|
||||
|
||||
# Check if only the inferred layout image was saved if extracted layout is None
|
||||
with (
|
||||
tempfile.TemporaryDirectory() as tmpdir,
|
||||
patch("unstructured.partition.pdf_image.pdf_image_utils.write_image") as mock_write_image,
|
||||
):
|
||||
pdf_image_utils.annotate_layout_elements_with_image(
|
||||
inferred_page_layout=inferred_layout,
|
||||
extracted_page_layout=None,
|
||||
output_dir_path=str(tmpdir),
|
||||
output_f_basename=output_basename,
|
||||
page_number=page_number,
|
||||
)
|
||||
|
||||
expected_filename = f"{output_basename}_{page_number}_inferred.jpg"
|
||||
actual_calls = [call.args[1] for call in mock_write_image.call_args_list]
|
||||
assert any(expected_filename in actual_call for actual_call in actual_calls)
|
||||
assert len(actual_calls) == 1 # Only one image should be saved
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "is_image"),
|
||||
[
|
||||
(example_doc_path("pdf/layout-parser-paper-fast.pdf"), False),
|
||||
(example_doc_path("img/layout-parser-paper-fast.jpg"), True),
|
||||
],
|
||||
)
|
||||
def test_annotate_layout_elements(filename, is_image):
|
||||
inferred_document_layout = MockDocumentLayout
|
||||
extracted_layout = [MagicMock(), MagicMock()]
|
||||
|
||||
with (
|
||||
patch("PIL.Image.open"),
|
||||
patch(
|
||||
"unstructured.partition.pdf_image.pdf_image_utils.convert_pdf_to_image",
|
||||
return_value=["/path/to/image1.jpg", "/path/to/image2.jpg"],
|
||||
) as mock_pdf2image,
|
||||
patch(
|
||||
"unstructured.partition.pdf_image.pdf_image_utils.annotate_layout_elements_with_image"
|
||||
) as mock_annotate_layout_elements_with_image,
|
||||
):
|
||||
pdf_image_utils.annotate_layout_elements(
|
||||
inferred_document_layout=inferred_document_layout,
|
||||
extracted_layout=extracted_layout,
|
||||
filename=filename,
|
||||
output_dir_path="/output",
|
||||
pdf_image_dpi=200,
|
||||
is_image=is_image,
|
||||
)
|
||||
if is_image:
|
||||
mock_annotate_layout_elements_with_image.assert_called_once()
|
||||
else:
|
||||
assert mock_annotate_layout_elements_with_image.call_count == len(
|
||||
mock_pdf2image.return_value
|
||||
)
|
||||
|
||||
|
||||
def test_annotate_layout_elements_file_not_found_error():
|
||||
with pytest.raises(FileNotFoundError):
|
||||
pdf_image_utils.annotate_layout_elements(
|
||||
inferred_document_layout=MagicMock(),
|
||||
extracted_layout=[],
|
||||
filename="nonexistent.jpg",
|
||||
output_dir_path="/output",
|
||||
pdf_image_dpi=200,
|
||||
is_image=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[("test\tco\x0cn\ftrol\ncharacter\rs\b", "test control characters"), ("\"'\\", "\"'\\")],
|
||||
)
|
||||
def test_remove_control_characters(text, expected):
|
||||
assert pdf_image_utils.remove_control_characters(text) == expected
|
||||
@@ -0,0 +1,368 @@
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from pdfminer.layout import LAParams, LTChar, LTContainer
|
||||
from PIL import Image
|
||||
from unstructured_inference.constants import IsExtracted
|
||||
from unstructured_inference.constants import Source as InferenceSource
|
||||
from unstructured_inference.inference.elements import (
|
||||
EmbeddedTextRegion,
|
||||
Rectangle,
|
||||
TextRegion,
|
||||
TextRegions,
|
||||
)
|
||||
from unstructured_inference.inference.layout import DocumentLayout, LayoutElement, PageLayout
|
||||
from unstructured_inference.inference.layoutelement import LayoutElements
|
||||
|
||||
from test_unstructured.unit_utils import example_doc_path
|
||||
from unstructured.partition.auto import partition
|
||||
from unstructured.partition.pdf_image.pdfminer_processing import (
|
||||
_validate_bbox,
|
||||
aggregate_embedded_text_by_block,
|
||||
bboxes1_is_almost_subregion_of_bboxes2,
|
||||
boxes_self_iou,
|
||||
clean_pdfminer_inner_elements,
|
||||
process_file_with_pdfminer,
|
||||
remove_duplicate_elements,
|
||||
text_is_embedded,
|
||||
)
|
||||
from unstructured.partition.utils.constants import Source
|
||||
|
||||
# A set of elements with pdfminer elements inside tables
|
||||
deletable_elements_inside_table = [
|
||||
LayoutElement(
|
||||
bbox=Rectangle(0, 0, 100, 100),
|
||||
text="Table with inner elements",
|
||||
type="Table",
|
||||
),
|
||||
LayoutElement(bbox=Rectangle(50, 50, 70, 70), text="text1", source=Source.PDFMINER),
|
||||
LayoutElement(bbox=Rectangle(70, 70, 80, 80), text="text2", source=Source.PDFMINER),
|
||||
]
|
||||
|
||||
# A set of elements without pdfminer elements inside
|
||||
# tables (no elements with source=Source.PDFMINER)
|
||||
no_deletable_elements_inside_table = [
|
||||
LayoutElement(
|
||||
bbox=Rectangle(0, 0, 100, 100),
|
||||
text="Table with inner elements",
|
||||
type="Table",
|
||||
source=InferenceSource.YOLOX,
|
||||
),
|
||||
LayoutElement(bbox=Rectangle(50, 50, 70, 70), text="text1", source=InferenceSource.YOLOX),
|
||||
LayoutElement(bbox=Rectangle(70, 70, 80, 80), text="text2", source=InferenceSource.YOLOX),
|
||||
]
|
||||
# A set of elements with pdfminer elements inside tables and other
|
||||
# elements with source=Source.PDFMINER
|
||||
# Note: there is some elements with source=Source.PDFMINER are not inside tables
|
||||
mix_elements_inside_table = [
|
||||
LayoutElement(
|
||||
bbox=Rectangle(0, 0, 100, 100),
|
||||
text="Table1 with inner elements",
|
||||
type="Table",
|
||||
source=InferenceSource.YOLOX,
|
||||
),
|
||||
LayoutElement(bbox=Rectangle(50, 50, 70, 70), text="Inside table1"),
|
||||
LayoutElement(bbox=Rectangle(70, 70, 80, 80), text="Inside table1", source=Source.PDFMINER),
|
||||
LayoutElement(
|
||||
bbox=Rectangle(150, 150, 170, 170),
|
||||
text="Outside tables",
|
||||
source=Source.PDFMINER,
|
||||
),
|
||||
LayoutElement(
|
||||
bbox=Rectangle(180, 180, 200, 200),
|
||||
text="Outside tables",
|
||||
source=Source.PDFMINER,
|
||||
),
|
||||
LayoutElement(
|
||||
bbox=Rectangle(0, 500, 100, 700),
|
||||
text="Table2 with inner elements",
|
||||
type="Table",
|
||||
source=InferenceSource.YOLOX,
|
||||
),
|
||||
LayoutElement(bbox=Rectangle(0, 510, 50, 600), text="Inside table2", source=Source.PDFMINER),
|
||||
LayoutElement(bbox=Rectangle(0, 550, 70, 650), text="Inside table2", source=Source.PDFMINER),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("bbox", "is_valid"),
|
||||
[
|
||||
([0, 1, 0, 1], False),
|
||||
([0, 1, 1, 2], True),
|
||||
([0, 1, 1, None], False),
|
||||
([0, 1, 1, np.nan], False),
|
||||
([0, 1, -1, 0], False),
|
||||
([0, 1, -1, 2], False),
|
||||
],
|
||||
)
|
||||
def test_valid_bbox(bbox, is_valid):
|
||||
assert _validate_bbox(bbox) is is_valid
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("elements", "length_extra_info", "expected_document_length"),
|
||||
[
|
||||
(deletable_elements_inside_table, 1, 1),
|
||||
(no_deletable_elements_inside_table, 0, 3),
|
||||
(mix_elements_inside_table, 2, 5),
|
||||
],
|
||||
)
|
||||
def test_clean_pdfminer_inner_elements(elements, length_extra_info, expected_document_length):
|
||||
# create a sample document with pdfminer elements inside tables
|
||||
page = PageLayout(number=1, image=Image.new("1", (1, 1)))
|
||||
page.elements_array = LayoutElements.from_list(elements)
|
||||
document_with_table = DocumentLayout(pages=[page])
|
||||
document = document_with_table
|
||||
|
||||
# call the function to clean the pdfminer inner elements
|
||||
cleaned_doc = clean_pdfminer_inner_elements(document)
|
||||
|
||||
# check that the pdfminer elements were stored in the extra_info dictionary
|
||||
assert len(cleaned_doc.pages[0].elements_array) == expected_document_length
|
||||
|
||||
|
||||
elements_with_duplicate_images = [
|
||||
LayoutElement(
|
||||
bbox=Rectangle(0, 0, 100, 100),
|
||||
text="Image1",
|
||||
type="Image",
|
||||
source=Source.PDFMINER,
|
||||
),
|
||||
LayoutElement(
|
||||
bbox=Rectangle(10, 10, 110, 110), text="Image1", type="Image", source=Source.PDFMINER
|
||||
),
|
||||
LayoutElement(bbox=Rectangle(150, 150, 170, 170), text="Title1", type="Title"),
|
||||
]
|
||||
|
||||
elements_without_duplicate_images = [
|
||||
LayoutElement(
|
||||
bbox=Rectangle(0, 0, 100, 100),
|
||||
text="Sample image",
|
||||
type="Image",
|
||||
source=Source.PDFMINER,
|
||||
),
|
||||
LayoutElement(
|
||||
bbox=Rectangle(10, 10, 110, 110),
|
||||
text="Sample image with similar bbox",
|
||||
type="Image",
|
||||
source=Source.PDFMINER,
|
||||
),
|
||||
LayoutElement(
|
||||
bbox=Rectangle(200, 200, 250, 250),
|
||||
text="Sample image",
|
||||
type="Image",
|
||||
source=Source.PDFMINER,
|
||||
),
|
||||
LayoutElement(bbox=Rectangle(150, 150, 170, 170), text="Title1", type="Title"),
|
||||
]
|
||||
|
||||
|
||||
def test_aggregate_by_block():
|
||||
expected = "Inside region1 Inside region2"
|
||||
embedded_regions = TextRegions.from_list(
|
||||
[
|
||||
TextRegion.from_coords(0, 0, 300, 20, "Inside region1"),
|
||||
TextRegion.from_coords(0, 20, 300, 80, None),
|
||||
TextRegion.from_coords(0, 80, 200, 300, "Inside region2"),
|
||||
TextRegion.from_coords(250, 250, 350, 350, "Outside region"),
|
||||
]
|
||||
)
|
||||
embedded_regions.is_extracted_array = np.array([IsExtracted.TRUE] * 4)
|
||||
target_region = TextRegions.from_list([TextRegion.from_coords(0, 0, 300, 300)])
|
||||
|
||||
text, extracted = aggregate_embedded_text_by_block(target_region, embedded_regions)
|
||||
assert text == expected
|
||||
assert extracted.value == "true"
|
||||
|
||||
|
||||
def test_aggregate_only_partially_fill_target():
|
||||
expected = "Inside region1"
|
||||
embedded_regions = TextRegions.from_list(
|
||||
[
|
||||
TextRegion.from_coords(0, 0, 20, 20, "Inside region1"),
|
||||
]
|
||||
)
|
||||
embedded_regions.is_extracted_array = np.array([IsExtracted.TRUE])
|
||||
target_region = TextRegions.from_list([TextRegion.from_coords(0, 0, 300, 300)])
|
||||
|
||||
text, extracted = aggregate_embedded_text_by_block(target_region, embedded_regions)
|
||||
assert text == expected
|
||||
assert extracted.value == "partial"
|
||||
|
||||
|
||||
def test_aggregate_not_filling_target():
|
||||
embedded_regions = TextRegions.from_list(
|
||||
[
|
||||
TextRegion.from_coords(300, 0, 400, 20, "outside"),
|
||||
]
|
||||
)
|
||||
embedded_regions.is_extracted_array = np.array([IsExtracted.TRUE])
|
||||
target_region = TextRegions.from_list([TextRegion.from_coords(0, 0, 300, 300)])
|
||||
|
||||
text, extracted = aggregate_embedded_text_by_block(target_region, embedded_regions)
|
||||
assert text == ""
|
||||
assert extracted.value == "false"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("coords1", "coords2", "expected"),
|
||||
[
|
||||
(
|
||||
[[0, 0, 10, 10], [10, 0, 20, 10], [10, 10, 20, 20]],
|
||||
[[0, 0, 10, 10], [0, 0, 12, 12]],
|
||||
[[True, True], [False, False], [False, False]],
|
||||
),
|
||||
(
|
||||
[[0, 0, 10, 10], [10, 0, 20, 10], [10, 10, 20, 20]],
|
||||
[[0, 0, 10, 10], [10, 10, 22, 22], [0, 0, 5, 5]],
|
||||
[[True, False, False], [False, False, False], [False, True, False]],
|
||||
),
|
||||
(
|
||||
[[0, 0, 10, 10], [10, 10, 10, 10]],
|
||||
[[0, 0, 10, 10], [10, 10, 22, 22], [0, 0, 5, 5]],
|
||||
[[True, False, False], [True, True, False]],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_bboxes1_is_almost_subregion_of_bboxes2(coords1, coords2, expected):
|
||||
bboxes1 = [Rectangle(*row) for row in coords1]
|
||||
bboxes2 = [Rectangle(*row) for row in coords2]
|
||||
np.testing.assert_array_equal(
|
||||
bboxes1_is_almost_subregion_of_bboxes2(bboxes1, bboxes2), expected
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("coords", "threshold", "expected"),
|
||||
[
|
||||
(
|
||||
[[0, 0, 10, 10], [2, 2, 12, 12], [10, 10, 20, 20]],
|
||||
0.5,
|
||||
[[True, True, False], [True, True, False], [False, False, True]],
|
||||
),
|
||||
(
|
||||
[[0, 0, 10, 10], [2, 2, 12, 12], [10, 10, 20, 20]],
|
||||
0.9,
|
||||
[[True, False, False], [False, True, False], [False, False, True]],
|
||||
),
|
||||
(
|
||||
[[0, 0, 10, 10], [10, 10, 10, 10]],
|
||||
0.5,
|
||||
[[True, False], [False, True]],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_boxes_self_iou(coords, threshold, expected):
|
||||
bboxes = [Rectangle(*row) for row in coords]
|
||||
np.testing.assert_array_equal(boxes_self_iou(bboxes, threshold), expected)
|
||||
|
||||
|
||||
def test_remove_duplicate_elements():
|
||||
sample_elements = TextRegions.from_list(
|
||||
[
|
||||
EmbeddedTextRegion(bbox=Rectangle(0, 0, 10, 10), text="Text 1"),
|
||||
EmbeddedTextRegion(bbox=Rectangle(0, 0, 10, 10), text="Text 2"),
|
||||
EmbeddedTextRegion(bbox=Rectangle(20, 20, 30, 30), text="Text 3"),
|
||||
]
|
||||
)
|
||||
|
||||
result = remove_duplicate_elements(sample_elements)
|
||||
|
||||
# Check that duplicates were removed and only 2 unique elements remain
|
||||
assert len(result) == 2
|
||||
assert result.texts.tolist() == ["Text 2", "Text 3"]
|
||||
assert result.element_coords.tolist() == [[0, 0, 10, 10], [20, 20, 30, 30]]
|
||||
|
||||
|
||||
def test_process_file_with_pdfminer():
|
||||
layout, links = process_file_with_pdfminer(example_doc_path("pdf/layout-parser-paper-fast.pdf"))
|
||||
assert len(layout)
|
||||
assert "LayoutParser: A Unified Toolkit for Deep\n" in layout[0].texts
|
||||
assert links[0][0]["url"] == "https://layout-parser.github.io"
|
||||
|
||||
|
||||
def test_process_file_with_pdfminer_is_extracted_array():
|
||||
layout, _ = process_file_with_pdfminer(example_doc_path("pdf/layout-parser-paper-fast.pdf"))
|
||||
assert all(is_extracted is IsExtracted.TRUE for is_extracted in layout[0].is_extracted_array)
|
||||
|
||||
|
||||
def test_process_file_hidden_ocr_text():
|
||||
"""Test processing a PDF that contains hidden OCR text layer.
|
||||
|
||||
Note: pdfminer >= 20251230 fixed color state handling (PR #1140), which means
|
||||
invisible OCR text can no longer be detected via scolor/ncolor being None.
|
||||
The rendermode check also doesn't work as LTChar doesn't expose textstate.render.
|
||||
As a result, all text is now marked as IsExtracted.TRUE.
|
||||
"""
|
||||
layout, _ = process_file_with_pdfminer(example_doc_path("pdf/pdf-with-ocr-text.pdf"))
|
||||
# Only check text elements (class_id == 0); images (class_id == 1) always have None
|
||||
text_mask = layout[0].element_class_ids == 0
|
||||
text_is_extracted = layout[0].is_extracted_array[text_mask]
|
||||
assert all(is_extracted is IsExtracted.TRUE for is_extracted in text_is_extracted)
|
||||
|
||||
|
||||
@patch("unstructured.partition.pdf_image.pdfminer_utils.LAParams", return_value=LAParams())
|
||||
def test_laprams_are_passed_from_partition_to_pdfminer(pdfminer_mock):
|
||||
partition(
|
||||
filename=example_doc_path("pdf/layout-parser-paper-fast.pdf"),
|
||||
pdfminer_line_margin=1.123,
|
||||
pdfminer_char_margin=None,
|
||||
pdfminer_line_overlap=0.0123,
|
||||
pdfminer_word_margin=3.21,
|
||||
)
|
||||
assert pdfminer_mock.call_args.kwargs == {
|
||||
"line_margin": 1.123,
|
||||
"line_overlap": 0.0123,
|
||||
"word_margin": 3.21,
|
||||
}
|
||||
|
||||
|
||||
def create_mock_ltchar(text, invisible=False):
|
||||
"""Create a mock LTChar object"""
|
||||
|
||||
graphicstate = Mock()
|
||||
if invisible:
|
||||
graphicstate.scolor = None
|
||||
graphicstate.ncolor = None
|
||||
|
||||
char = LTChar(
|
||||
matrix=(1, 0, 0, 1, 0, 0), # transformation matrix
|
||||
font=Mock(), # you'd need to mock PDFFont
|
||||
fontsize=12,
|
||||
scaling=1,
|
||||
rise=0,
|
||||
text=text,
|
||||
textwidth=10,
|
||||
textdisp=(0, 1),
|
||||
ncs=Mock(),
|
||||
graphicstate=graphicstate,
|
||||
)
|
||||
|
||||
return char
|
||||
|
||||
|
||||
def create_mock_ltcontainer(chars):
|
||||
"""Create a mock LTContainer with LTChar objects"""
|
||||
container = LTContainer(bbox=(0, 0, 1, 1))
|
||||
|
||||
# The container should be iterable
|
||||
container.extend(chars)
|
||||
|
||||
return container
|
||||
|
||||
|
||||
# Now you can use it in tests
|
||||
def test_text_is_embedded():
|
||||
chars = [
|
||||
create_mock_ltchar("H"),
|
||||
create_mock_ltchar("e"),
|
||||
create_mock_ltchar("l"),
|
||||
create_mock_ltchar("l"),
|
||||
create_mock_ltchar("o", invisible=True),
|
||||
]
|
||||
|
||||
container = create_mock_ltcontainer(chars)
|
||||
|
||||
assert text_is_embedded(container, threshold=0.5)
|
||||
assert not text_is_embedded(container, threshold=0.1)
|
||||
@@ -0,0 +1,28 @@
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from pdfminer.layout import LTContainer, LTTextLine
|
||||
|
||||
from unstructured.partition.pdf_image.pdfminer_utils import extract_text_objects
|
||||
|
||||
|
||||
def test_extract_text_objects_nested_containers():
|
||||
"""Test extract_text_objects with nested LTContainers."""
|
||||
# Mock LTTextLine objects
|
||||
mock_text_line1 = MagicMock(spec=LTTextLine)
|
||||
mock_text_line2 = MagicMock(spec=LTTextLine)
|
||||
|
||||
# Mock inner container containing one LTTextLine
|
||||
mock_inner_container = MagicMock(spec=LTContainer)
|
||||
mock_inner_container.__iter__.return_value = [mock_text_line2]
|
||||
|
||||
# Mock outer container containing another LTTextLine and the inner container
|
||||
mock_outer_container = MagicMock(spec=LTContainer)
|
||||
mock_outer_container.__iter__.return_value = [mock_text_line1, mock_inner_container]
|
||||
|
||||
# Call the function with the outer container
|
||||
result = extract_text_objects(mock_outer_container)
|
||||
|
||||
# Assert both text line objects are extracted, even from nested containers
|
||||
assert len(result) == 2
|
||||
assert mock_text_line1 in result
|
||||
assert mock_text_line2 in result
|
||||
@@ -0,0 +1,660 @@
|
||||
import base64
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from unstructured_client.general import General
|
||||
from unstructured_client.models import shared
|
||||
from unstructured_client.models.operations import PartitionRequest
|
||||
from unstructured_client.models.shared import PartitionParameters
|
||||
from unstructured_client.utils import retries
|
||||
|
||||
from unstructured.documents.elements import ElementType, NarrativeText
|
||||
from unstructured.partition.api import (
|
||||
DEFAULT_RETRIES_MAX_ELAPSED_TIME_SEC,
|
||||
DEFAULT_RETRIES_MAX_INTERVAL_SEC,
|
||||
get_retries_config,
|
||||
partition_multiple_via_api,
|
||||
partition_via_api,
|
||||
)
|
||||
|
||||
from ..unit_utils import ANY, FixtureRequest, example_doc_path, method_mock
|
||||
|
||||
DIRECTORY = pathlib.Path(__file__).parent.resolve()
|
||||
|
||||
# NOTE(yao): point to paid API for now
|
||||
API_URL = "https://api.unstructuredapp.io/general/v0/general"
|
||||
|
||||
is_in_ci = os.getenv("CI", "").lower() not in {"", "false", "f", "0"}
|
||||
skip_not_on_main = os.getenv("GITHUB_REF_NAME", "").lower() != "main"
|
||||
|
||||
|
||||
def test_partition_via_api_with_filename_correctly_calls_sdk(
|
||||
request: FixtureRequest, expected_call_: list[Any]
|
||||
):
|
||||
partition_mock_ = method_mock(
|
||||
request, General, "partition", return_value=FakeResponse(status_code=200)
|
||||
)
|
||||
|
||||
elements = partition_via_api(filename=example_doc_path("eml/fake-email.eml"))
|
||||
|
||||
partition_mock_.assert_called_once_with(
|
||||
expected_call_[0], request=expected_call_[1], retries=expected_call_[2]
|
||||
)
|
||||
assert isinstance(partition_mock_.call_args_list[0].args[0], General)
|
||||
assert len(elements) == 1
|
||||
assert elements[0] == NarrativeText("This is a test email to use for unit tests.")
|
||||
assert elements[0].metadata.filetype == "message/rfc822"
|
||||
|
||||
|
||||
def test_partition_via_api_with_file_correctly_calls_sdk(
|
||||
request: FixtureRequest, expected_call_: list[Any]
|
||||
):
|
||||
partition_mock_ = method_mock(
|
||||
request, General, "partition", return_value=FakeResponse(status_code=200)
|
||||
)
|
||||
|
||||
with open(example_doc_path("eml/fake-email.eml"), "rb") as f:
|
||||
elements = partition_via_api(
|
||||
file=f, metadata_filename=example_doc_path("eml/fake-email.eml")
|
||||
)
|
||||
|
||||
# Update the fixture content to match the format passed to partition_via_api
|
||||
modified_expected_call = expected_call_[:]
|
||||
modified_expected_call[1].partition_parameters.files.content = f
|
||||
|
||||
partition_mock_.assert_called_once_with(
|
||||
modified_expected_call[0],
|
||||
request=modified_expected_call[1],
|
||||
retries=modified_expected_call[2],
|
||||
)
|
||||
assert isinstance(partition_mock_.call_args_list[0].args[0], General)
|
||||
assert len(elements) == 1
|
||||
assert elements[0] == NarrativeText("This is a test email to use for unit tests.")
|
||||
assert elements[0].metadata.filetype == "message/rfc822"
|
||||
|
||||
|
||||
def test_partition_via_api_warns_with_file_and_filename_and_calls_sdk(
|
||||
request: FixtureRequest, expected_call_: list[Any], caplog: pytest.LogCaptureFixture
|
||||
):
|
||||
partition_mock_ = method_mock(
|
||||
request, General, "partition", return_value=FakeResponse(status_code=200)
|
||||
)
|
||||
|
||||
with open(example_doc_path("eml/fake-email.eml"), "rb") as f:
|
||||
partition_via_api(file=f, file_filename=example_doc_path("eml/fake-email.eml"))
|
||||
|
||||
# Update the fixture content to match the format passed to partition_via_api
|
||||
modified_expected_call = expected_call_[:]
|
||||
modified_expected_call[1].partition_parameters.files.content = f
|
||||
|
||||
partition_mock_.assert_called_once_with(
|
||||
modified_expected_call[0],
|
||||
request=modified_expected_call[1],
|
||||
retries=modified_expected_call[2],
|
||||
)
|
||||
assert "WARNING" in caplog.text
|
||||
assert "The file_filename kwarg will be deprecated" in caplog.text
|
||||
|
||||
|
||||
def test_partition_via_api_from_file_raises_with_metadata_and_file_and_filename():
|
||||
filename = example_doc_path("eml/fake-email.eml")
|
||||
|
||||
with open(filename, "rb") as f, pytest.raises(ValueError):
|
||||
partition_via_api(file=f, file_filename=filename, metadata_filename=filename)
|
||||
|
||||
|
||||
def test_partition_via_api_from_file_raises_without_filename():
|
||||
with open(example_doc_path("eml/fake-email.eml"), "rb") as f, pytest.raises(ValueError):
|
||||
partition_via_api(file=f)
|
||||
|
||||
|
||||
def test_partition_via_api_raises_with_bad_response(request: FixtureRequest):
|
||||
partition_mock_ = method_mock(
|
||||
request, General, "partition", return_value=FakeResponse(status_code=500)
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_via_api(filename=example_doc_path("eml/fake-email.eml"))
|
||||
partition_mock_.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_in_ci, reason="Skipping test run outside of CI")
|
||||
@pytest.mark.skipif(skip_not_on_main, reason="Skipping test run outside of main branch")
|
||||
def test_partition_via_api_with_no_strategy():
|
||||
test_file = example_doc_path("pdf/loremipsum-flat.pdf")
|
||||
elements_no_strategy = partition_via_api(
|
||||
filename=test_file,
|
||||
strategy="auto",
|
||||
api_key=get_api_key(),
|
||||
# The url has changed since the 06/24 API release while the sdk defaults to the old url
|
||||
api_url=API_URL,
|
||||
skip_infer_table_types=["pdf"],
|
||||
)
|
||||
elements_hi_res = partition_via_api(
|
||||
filename=test_file,
|
||||
strategy="hi_res",
|
||||
api_key=get_api_key(),
|
||||
# The url has changed since the 06/24 API release while the sdk defaults to the old url
|
||||
api_url=API_URL,
|
||||
skip_infer_table_types=["pdf"],
|
||||
)
|
||||
elements_fast_res = partition_via_api(
|
||||
filename=test_file,
|
||||
strategy="fast",
|
||||
api_key=get_api_key(),
|
||||
# The url has changed since the 06/24 API release while the sdk defaults to the old url
|
||||
api_url=API_URL,
|
||||
skip_infer_table_types=["pdf"],
|
||||
)
|
||||
|
||||
# confirm that hi_res strategy was not passed as default to partition by comparing outputs
|
||||
# elements_hi_res[3].text =
|
||||
# 'LayoutParser: A Unified Toolkit for Deep Learning Based Document Image Analysis'
|
||||
# while elements_no_strategy[3].text = ']' (as of this writing)
|
||||
assert len(elements_no_strategy) == len(elements_hi_res)
|
||||
assert len(elements_hi_res) != len(elements_fast_res)
|
||||
|
||||
# NOTE(crag): slightly out scope assertion, but avoid extra API call
|
||||
assert elements_hi_res[0].metadata.coordinates is None
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_in_ci, reason="Skipping test run outside of CI")
|
||||
@pytest.mark.skipif(skip_not_on_main, reason="Skipping test run outside of main branch")
|
||||
def test_partition_via_api_with_image_hi_res_strategy_includes_coordinates():
|
||||
# coordinates not included by default to limit payload size
|
||||
elements = partition_via_api(
|
||||
filename=example_doc_path("pdf/fake-memo.pdf"),
|
||||
strategy="hi_res",
|
||||
coordinates="true",
|
||||
api_key=get_api_key(),
|
||||
api_url=API_URL,
|
||||
)
|
||||
|
||||
assert elements[0].metadata.coordinates is not None
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_in_ci, reason="Skipping test run outside of CI")
|
||||
@pytest.mark.skipif(skip_not_on_main, reason="Skipping test run outside of main branch")
|
||||
def test_partition_via_api_image_block_extraction():
|
||||
elements = partition_via_api(
|
||||
filename=example_doc_path("pdf/embedded-images-tables.pdf"),
|
||||
strategy="hi_res",
|
||||
extract_image_block_types=["image", "table"],
|
||||
api_key=get_api_key(),
|
||||
# The url has changed since the 06/24 API release while the sdk defaults to the old url
|
||||
api_url=API_URL,
|
||||
)
|
||||
image_elements = [el for el in elements if el.category == ElementType.IMAGE]
|
||||
for el in image_elements:
|
||||
assert el.metadata.image_base64 is not None
|
||||
assert el.metadata.image_mime_type is not None
|
||||
image_data = base64.b64decode(el.metadata.image_base64)
|
||||
assert isinstance(image_data, bytes)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_in_ci, reason="Skipping test run outside of CI")
|
||||
@pytest.mark.skipif(skip_not_on_main, reason="Skipping test run outside of main branch")
|
||||
def test_partition_via_api_retries_config():
|
||||
elements = partition_via_api(
|
||||
filename=example_doc_path("pdf/embedded-images-tables.pdf"),
|
||||
strategy="fast",
|
||||
api_key=get_api_key(),
|
||||
# The url has changed since the 06/24 API release while the sdk defaults to the old url
|
||||
api_url=API_URL,
|
||||
retries_initial_interval=5,
|
||||
retries_max_interval=15,
|
||||
retries_max_elapsed_time=100,
|
||||
retries_connection_errors=True,
|
||||
retries_exponent=1.5,
|
||||
)
|
||||
|
||||
assert len(elements) > 0
|
||||
|
||||
|
||||
# Note(austin) - This test is way too noisy against the hosted api
|
||||
# def test_partition_via_api_invalid_request_data_kwargs():
|
||||
# filename = os.path.join(DIRECTORY, "..", "..", "example-docs", "layout-parser-paper-fast.pdf")
|
||||
# with pytest.raises(SDKError):
|
||||
# partition_via_api(filename=filename, strategy="not_a_strategy")
|
||||
|
||||
|
||||
def test_retries_config_with_parameters_set():
|
||||
sdk = Mock()
|
||||
retries_config = get_retries_config(
|
||||
retries_connection_errors=True,
|
||||
retries_exponent=1.75,
|
||||
retries_initial_interval=20,
|
||||
retries_max_elapsed_time=1000,
|
||||
retries_max_interval=100,
|
||||
sdk=sdk,
|
||||
)
|
||||
|
||||
assert retries_config.retry_connection_errors
|
||||
assert retries_config.backoff.exponent == 1.75
|
||||
assert retries_config.backoff.initial_interval == 20
|
||||
assert retries_config.backoff.max_elapsed_time == 1000
|
||||
assert retries_config.backoff.max_interval == 100
|
||||
|
||||
|
||||
def test_retries_config_none_parameters_return_empty_config():
|
||||
sdk = Mock()
|
||||
retries_config = get_retries_config(
|
||||
retries_connection_errors=None,
|
||||
retries_exponent=None,
|
||||
retries_initial_interval=None,
|
||||
retries_max_elapsed_time=None,
|
||||
retries_max_interval=None,
|
||||
sdk=sdk,
|
||||
)
|
||||
|
||||
assert retries_config is None
|
||||
|
||||
|
||||
def test_retry_config_with_empty_sdk_retry_config_returns_default():
|
||||
sdk = Mock()
|
||||
sdk.sdk_configuration.retry_config = None
|
||||
retries_config = get_retries_config(
|
||||
retries_connection_errors=True,
|
||||
retries_exponent=1.88,
|
||||
retries_initial_interval=3000,
|
||||
retries_max_elapsed_time=None,
|
||||
retries_max_interval=None,
|
||||
sdk=sdk,
|
||||
)
|
||||
|
||||
assert retries_config.retry_connection_errors
|
||||
assert retries_config.backoff.exponent == 1.88
|
||||
assert retries_config.backoff.initial_interval == 3000
|
||||
assert retries_config.backoff.max_elapsed_time == DEFAULT_RETRIES_MAX_ELAPSED_TIME_SEC
|
||||
assert retries_config.backoff.max_interval == DEFAULT_RETRIES_MAX_INTERVAL_SEC
|
||||
|
||||
|
||||
def test_retries_config_with_no_parameters_set():
|
||||
retry_config = retries.RetryConfig(
|
||||
"backoff", retries.BackoffStrategy(3000, 720000, 1.88, 1800000), True
|
||||
)
|
||||
sdk = Mock()
|
||||
sdk.sdk_configuration.retry_config = retry_config
|
||||
retries_config = get_retries_config(
|
||||
retries_connection_errors=True,
|
||||
retries_exponent=None,
|
||||
retries_initial_interval=None,
|
||||
retries_max_elapsed_time=None,
|
||||
retries_max_interval=None,
|
||||
sdk=sdk,
|
||||
)
|
||||
|
||||
assert retries_config.retry_connection_errors
|
||||
assert retries_config.backoff.exponent == 1.88
|
||||
assert retries_config.backoff.initial_interval == 3000
|
||||
assert retries_config.backoff.max_elapsed_time == 1800000
|
||||
assert retries_config.backoff.max_interval == 720000
|
||||
|
||||
|
||||
def test_retries_config_cascade():
|
||||
# notice max_interval is set to 0 which is incorrect - so the DEFAULT_RETRIES_MAX_INTERVAL_SEC
|
||||
# should be used
|
||||
retry_config = retries.RetryConfig(
|
||||
"backoff", retries.BackoffStrategy(3000, 0, 1.88, None), True
|
||||
)
|
||||
sdk = Mock()
|
||||
sdk.sdk_configuration.retry_config = retry_config
|
||||
retries_config = get_retries_config(
|
||||
retries_connection_errors=False,
|
||||
retries_exponent=1.75,
|
||||
retries_initial_interval=20,
|
||||
retries_max_elapsed_time=None,
|
||||
retries_max_interval=None,
|
||||
sdk=sdk,
|
||||
)
|
||||
|
||||
assert not retries_config.retry_connection_errors
|
||||
assert retries_config.backoff.exponent == 1.75
|
||||
assert retries_config.backoff.initial_interval == 20
|
||||
assert retries_config.backoff.max_elapsed_time == DEFAULT_RETRIES_MAX_ELAPSED_TIME_SEC
|
||||
assert retries_config.backoff.max_interval == DEFAULT_RETRIES_MAX_INTERVAL_SEC
|
||||
|
||||
|
||||
def test_partition_multiple_via_api_with_single_filename(request: FixtureRequest):
|
||||
partition_mock_ = method_mock(
|
||||
request, requests, "post", return_value=FakeResponse(status_code=200)
|
||||
)
|
||||
filename = example_doc_path("eml/fake-email.eml")
|
||||
|
||||
elements = partition_multiple_via_api(filenames=[filename])
|
||||
|
||||
partition_mock_.assert_called_once_with(
|
||||
"https://api.unstructured.io/general/v0/general",
|
||||
headers={"ACCEPT": "application/json", "UNSTRUCTURED-API-KEY": ANY},
|
||||
data={},
|
||||
files=[("files", (example_doc_path("eml/fake-email.eml"), ANY, None))],
|
||||
)
|
||||
assert elements[0][0] == NarrativeText("This is a test email to use for unit tests.")
|
||||
assert elements[0][0].metadata.filetype == "message/rfc822"
|
||||
|
||||
|
||||
def test_partition_multiple_via_api_from_filenames(request: FixtureRequest):
|
||||
partition_mock_ = method_mock(
|
||||
request, requests, "post", return_value=FakeMultipleResponse(status_code=200)
|
||||
)
|
||||
filenames = [example_doc_path("eml/fake-email.eml"), example_doc_path("fake.docx")]
|
||||
|
||||
elements = partition_multiple_via_api(filenames=filenames)
|
||||
|
||||
partition_mock_.assert_called_once_with(
|
||||
"https://api.unstructured.io/general/v0/general",
|
||||
headers={"ACCEPT": "application/json", "UNSTRUCTURED-API-KEY": ANY},
|
||||
data={},
|
||||
files=[
|
||||
("files", (example_doc_path("eml/fake-email.eml"), ANY, None)),
|
||||
("files", (example_doc_path("fake.docx"), ANY, None)),
|
||||
],
|
||||
)
|
||||
assert len(elements) == 2
|
||||
assert elements[0][0] == NarrativeText("This is a test email to use for unit tests.")
|
||||
assert elements[0][0].metadata.filetype == "message/rfc822"
|
||||
|
||||
|
||||
def test_partition_multiple_via_api_from_files(request: FixtureRequest):
|
||||
partition_mock_ = method_mock(
|
||||
request, requests, "post", return_value=FakeMultipleResponse(status_code=200)
|
||||
)
|
||||
filenames = [example_doc_path("eml/fake-email.eml"), example_doc_path("fake.docx")]
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
files = [stack.enter_context(open(filename, "rb")) for filename in filenames]
|
||||
elements = partition_multiple_via_api(
|
||||
files=files,
|
||||
metadata_filenames=filenames,
|
||||
)
|
||||
|
||||
partition_mock_.assert_called_once_with(
|
||||
"https://api.unstructured.io/general/v0/general",
|
||||
headers={"ACCEPT": "application/json", "UNSTRUCTURED-API-KEY": ANY},
|
||||
data={},
|
||||
files=[
|
||||
("files", (example_doc_path("eml/fake-email.eml"), ANY, None)),
|
||||
("files", (example_doc_path("fake.docx"), ANY, None)),
|
||||
],
|
||||
)
|
||||
assert len(elements) == 2
|
||||
assert elements[0][0] == NarrativeText("This is a test email to use for unit tests.")
|
||||
assert elements[0][0].metadata.filetype == "message/rfc822"
|
||||
|
||||
|
||||
def test_partition_multiple_via_api_warns_with_file_filename(
|
||||
caplog: pytest.LogCaptureFixture, request: FixtureRequest
|
||||
):
|
||||
partition_mock_ = method_mock(
|
||||
request, requests, "post", return_value=FakeMultipleResponse(status_code=200)
|
||||
)
|
||||
filenames = [example_doc_path("eml/fake-email.eml"), example_doc_path("fake.docx")]
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
files = [stack.enter_context(open(filename, "rb")) for filename in filenames]
|
||||
partition_multiple_via_api(
|
||||
files=files,
|
||||
file_filenames=filenames,
|
||||
)
|
||||
|
||||
partition_mock_.assert_called_once_with(
|
||||
"https://api.unstructured.io/general/v0/general",
|
||||
headers={"ACCEPT": "application/json", "UNSTRUCTURED-API-KEY": ANY},
|
||||
data={},
|
||||
files=[
|
||||
("files", (example_doc_path("eml/fake-email.eml"), ANY, None)),
|
||||
("files", (example_doc_path("fake.docx"), ANY, None)),
|
||||
],
|
||||
)
|
||||
assert "WARNING" in caplog.text
|
||||
assert "The file_filenames kwarg will be deprecated" in caplog.text
|
||||
|
||||
|
||||
def test_partition_multiple_via_api_raises_with_file_and_metadata_filename():
|
||||
filenames = [example_doc_path("eml/fake-email.eml"), example_doc_path("fake.docx")]
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
files = [stack.enter_context(open(filename, "rb")) for filename in filenames]
|
||||
with pytest.raises(ValueError):
|
||||
partition_multiple_via_api(
|
||||
files=files,
|
||||
metadata_filenames=filenames,
|
||||
file_filenames=filenames,
|
||||
)
|
||||
|
||||
|
||||
def test_partition_multiple_via_api_raises_with_bad_response(request: FixtureRequest):
|
||||
partition_mock_ = method_mock(
|
||||
request, requests, "post", return_value=FakeMultipleResponse(status_code=500)
|
||||
)
|
||||
filenames = [example_doc_path("eml/fake-email.eml"), example_doc_path("fake.docx")]
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_multiple_via_api(filenames=filenames)
|
||||
partition_mock_.assert_called_once_with(
|
||||
"https://api.unstructured.io/general/v0/general",
|
||||
headers={"ACCEPT": "application/json", "UNSTRUCTURED-API-KEY": ANY},
|
||||
data={},
|
||||
files=[
|
||||
("files", (example_doc_path("eml/fake-email.eml"), ANY, None)),
|
||||
("files", (example_doc_path("fake.docx"), ANY, None)),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_partition_multiple_via_api_raises_with_content_types_size_mismatch():
|
||||
filenames = [example_doc_path("eml/fake-email.eml"), example_doc_path("fake.docx")]
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_multiple_via_api(
|
||||
filenames=filenames,
|
||||
content_types=["text/plain"],
|
||||
)
|
||||
|
||||
|
||||
def test_partition_multiple_via_api_from_files_raises_with_size_mismatch():
|
||||
filenames = [example_doc_path("eml/fake-email.eml"), example_doc_path("fake.docx")]
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
files = [stack.enter_context(open(filename, "rb")) for filename in filenames]
|
||||
with pytest.raises(ValueError):
|
||||
partition_multiple_via_api(
|
||||
files=files,
|
||||
metadata_filenames=filenames,
|
||||
content_types=["text/plain"],
|
||||
)
|
||||
|
||||
|
||||
def test_partition_multiple_via_api_from_files_raises_without_filenames():
|
||||
filenames = [example_doc_path("eml/fake-email.eml"), example_doc_path("fake.docx")]
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
files = [stack.enter_context(open(filename, "rb")) for filename in filenames]
|
||||
with pytest.raises(ValueError):
|
||||
partition_multiple_via_api(
|
||||
files=files,
|
||||
)
|
||||
|
||||
|
||||
def get_api_key():
|
||||
api_key = os.getenv("UNS_API_KEY")
|
||||
if api_key is None:
|
||||
raise ValueError("UNS_API_KEY environment variable not set")
|
||||
return api_key
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_in_ci, reason="Skipping test run outside of CI")
|
||||
@pytest.mark.skipif(skip_not_on_main, reason="Skipping test run outside of main branch")
|
||||
def test_partition_multiple_via_api_valid_request_data_kwargs():
|
||||
filenames = [
|
||||
example_doc_path("fake-text.txt"),
|
||||
example_doc_path("fake-email.txt"),
|
||||
]
|
||||
|
||||
list_of_lists_of_elements = partition_multiple_via_api(
|
||||
filenames=filenames,
|
||||
strategy="fast",
|
||||
api_key=get_api_key(),
|
||||
api_url=API_URL,
|
||||
)
|
||||
# assert there is a list of elements for each file
|
||||
assert len(list_of_lists_of_elements) == 2
|
||||
assert isinstance(list_of_lists_of_elements[0], list)
|
||||
assert isinstance(list_of_lists_of_elements[1], list)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_in_ci, reason="Skipping test run outside of CI")
|
||||
def test_partition_multiple_via_api_invalid_request_data_kwargs():
|
||||
filenames = [
|
||||
example_doc_path("pdf/layout-parser-paper-fast.pdf"),
|
||||
example_doc_path("img/layout-parser-paper-fast.jpg"),
|
||||
]
|
||||
with pytest.raises(ValueError):
|
||||
partition_multiple_via_api(
|
||||
filenames=filenames,
|
||||
strategy="not_a_strategy",
|
||||
api_key=get_api_key(),
|
||||
# The url has changed since the 06/24 API release while the sdk defaults to the old url
|
||||
api_url=API_URL,
|
||||
)
|
||||
|
||||
|
||||
MOCK_TEXT = """[
|
||||
{
|
||||
"element_id": "f49fbd614ddf5b72e06f59e554e6ae2b",
|
||||
"text": "This is a test email to use for unit tests.",
|
||||
"type": "NarrativeText",
|
||||
"metadata": {
|
||||
"sent_from": [
|
||||
"Matthew Robinson <mrobinson@unstructured.io>"
|
||||
],
|
||||
"sent_to": [
|
||||
"Matthew Robinson <mrobinson@unstructured.io>"
|
||||
],
|
||||
"subject": "Test Email",
|
||||
"filename": "fake-email.eml",
|
||||
"filetype": "message/rfc822"
|
||||
}
|
||||
}
|
||||
]"""
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, status_code: int):
|
||||
self.status_code = status_code
|
||||
# The string representation of partitioned elements is nested in an additional
|
||||
# layer in the new unstructured-client:
|
||||
# `elements_from_json(text=response.raw_response.text)`
|
||||
self.raw_response = FakeRawResponse()
|
||||
self.headers = {"Content-Type": "application/json"}
|
||||
|
||||
def json(self):
|
||||
return json.loads(self.text)
|
||||
|
||||
@property
|
||||
def text(self):
|
||||
return MOCK_TEXT
|
||||
|
||||
|
||||
class FakeRawResponse:
|
||||
def __init__(self):
|
||||
self.text = MOCK_TEXT
|
||||
|
||||
|
||||
class FakeMultipleResponse:
|
||||
def __init__(self, status_code: int):
|
||||
self.status_code = status_code
|
||||
|
||||
def json(self):
|
||||
return json.loads(self.text)
|
||||
|
||||
@property
|
||||
def text(self):
|
||||
return """[
|
||||
[
|
||||
{
|
||||
"element_id": "f49fbd614ddf5b72e06f59e554e6ae2b",
|
||||
"text": "This is a test email to use for unit tests.",
|
||||
"type": "NarrativeText",
|
||||
"metadata": {
|
||||
"sent_from": [
|
||||
"Matthew Robinson <mrobinson@unstructured.io>"
|
||||
],
|
||||
"sent_to": [
|
||||
"Matthew Robinson <mrobinson@unstructured.io>"
|
||||
],
|
||||
"subject": "Test Email",
|
||||
"filename": "fake-email.eml",
|
||||
"filetype": "message/rfc822"
|
||||
}
|
||||
}
|
||||
],
|
||||
[
|
||||
{
|
||||
"element_id": "f49fbd614ddf5b72e06f59e554e6ae2b",
|
||||
"text": "This is a test email to use for unit tests.",
|
||||
"type": "NarrativeText",
|
||||
"metadata": {
|
||||
"sent_from": [
|
||||
"Matthew Robinson <mrobinson@unstructured.io>"
|
||||
],
|
||||
"sent_to": [
|
||||
"Matthew Robinson <mrobinson@unstructured.io>"
|
||||
],
|
||||
"subject": "Test Email",
|
||||
"filename": "fake-email.eml",
|
||||
"filetype": "message/rfc822"
|
||||
}
|
||||
}
|
||||
]
|
||||
]"""
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def expected_call_():
|
||||
with open(example_doc_path("eml/fake-email.eml"), "rb") as f:
|
||||
file_bytes = f.read()
|
||||
return [
|
||||
ANY,
|
||||
PartitionRequest(
|
||||
partition_parameters=PartitionParameters(
|
||||
files=shared.Files(
|
||||
content=file_bytes,
|
||||
file_name=example_doc_path("eml/fake-email.eml"),
|
||||
),
|
||||
chunking_strategy=None,
|
||||
combine_under_n_chars=None,
|
||||
coordinates=False,
|
||||
encoding=None,
|
||||
extract_image_block_types=None,
|
||||
gz_uncompressed_content_type=None,
|
||||
hi_res_model_name=None,
|
||||
include_orig_elements=None,
|
||||
include_page_breaks=False,
|
||||
languages=None,
|
||||
max_characters=None,
|
||||
multipage_sections=True,
|
||||
new_after_n_chars=None,
|
||||
ocr_languages=None,
|
||||
output_format=shared.OutputFormat.APPLICATION_JSON,
|
||||
overlap=0,
|
||||
overlap_all=False,
|
||||
pdf_infer_table_structure=True,
|
||||
similarity_threshold=None,
|
||||
skip_infer_table_types=None,
|
||||
split_pdf_concurrency_level=5,
|
||||
split_pdf_page=True,
|
||||
starting_page_number=None,
|
||||
strategy=shared.Strategy.HI_RES,
|
||||
unique_element_ids=False,
|
||||
xml_keep_tags=False,
|
||||
)
|
||||
),
|
||||
None, # retries kwarg
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,142 @@
|
||||
EXPECTED_TABLE = (
|
||||
"<table>"
|
||||
"<tr><td>Stanley Cups</td><td/><td/></tr>"
|
||||
"<tr><td>Team</td><td>Location</td><td>Stanley Cups</td></tr>"
|
||||
"<tr><td>Blues</td><td>STL</td><td>1</td></tr>"
|
||||
"<tr><td>Flyers</td><td>PHI</td><td>2</td></tr>"
|
||||
"<tr><td>Maple Leafs</td><td>TOR</td><td>13</td></tr>"
|
||||
"</table>"
|
||||
)
|
||||
|
||||
EXPECTED_TABLE_SEMICOLON_DELIMITER = (
|
||||
"<table>"
|
||||
"<tr><td>Year</td><td>Month</td><td>Revenue</td><td>Costs</td><td/></tr>"
|
||||
"<tr><td>2022</td><td>1</td><td>123</td><td>-123</td><td/></tr>"
|
||||
"<tr><td>2023</td><td>2</td><td>143,1</td><td>-814,38</td><td/></tr>"
|
||||
"<tr><td>2024</td><td>3</td><td>215,32</td><td>-11,08</td><td/></tr>"
|
||||
"</table>"
|
||||
)
|
||||
|
||||
EXPECTED_TABLE_WITH_EMOJI = (
|
||||
"<table>"
|
||||
"<tr><td>Stanley Cups</td><td/><td/></tr>"
|
||||
"<tr><td>Team</td><td>Location</td><td>Stanley Cups</td></tr>"
|
||||
"<tr><td>Blues</td><td>STL</td><td>1</td></tr>"
|
||||
"<tr><td>Flyers</td><td>PHI</td><td>2</td></tr>"
|
||||
"<tr><td>Maple Leafs</td><td>TOR</td><td>13</td></tr>"
|
||||
"<tr><td>👨\\U+1F3FB🔧</td><td>TOR</td><td>15</td></tr>"
|
||||
"</table>"
|
||||
)
|
||||
|
||||
EXPECTED_TABLE_XLSX = (
|
||||
"<table>"
|
||||
"<tr><td>Team</td><td>Location</td><td>Stanley Cups</td></tr>"
|
||||
"<tr><td>Blues</td><td>STL</td><td>1</td></tr>"
|
||||
"<tr><td>Flyers</td><td>PHI</td><td>2</td></tr>"
|
||||
"<tr><td>Maple Leafs</td><td>TOR</td><td>13</td></tr>"
|
||||
"</table>"
|
||||
)
|
||||
|
||||
EXPECTED_TABLE_WITH_LINE_DELIMITER = (
|
||||
"<table>"
|
||||
"<tr><td>col1</td><td>col2</td><td>col3</td></tr>"
|
||||
"<tr><td>a</td><td>b</td><td>c</td></tr>"
|
||||
"<tr><td>d</td><td>e</td><td>f</td></tr>"
|
||||
"<tr><td>g</td><td>h</td><td>i</td></tr>"
|
||||
"</table>"
|
||||
)
|
||||
|
||||
|
||||
EXPECTED_TITLE = "Stanley Cups"
|
||||
|
||||
EXPECTED_TEXT = (
|
||||
"Stanley Cups Team Location Stanley Cups Blues STL 1 Flyers PHI 2 Maple Leafs TOR 13"
|
||||
)
|
||||
|
||||
EXPECTED_TEXT_XLSX = "Team Location Stanley Cups Blues STL 1 Flyers PHI 2 Maple Leafs TOR 13"
|
||||
|
||||
EXPECTED_TEXT_WITH_EMOJI = (
|
||||
"Stanley Cups "
|
||||
"Team Location Stanley Cups Blues STL 1 Flyers PHI 2 Maple Leafs TOR 13 👨\\U+1F3FB🔧 TOR 15"
|
||||
)
|
||||
|
||||
EXPECTED_TEXT_SEMICOLON_DELIMITER = (
|
||||
"Year Month Revenue Costs 2022 1 123 -123 2023 2 143,1 -814,38 2024 3 215,32 -11,08"
|
||||
)
|
||||
|
||||
EXPECTED_TEXT_WITH_LINE_DELIMITER = "col1 col2 col3 a b c d e f g h i"
|
||||
|
||||
EXPECTED_XLS_TABLE = (
|
||||
"<table><tr>"
|
||||
"<td>MC</td>"
|
||||
"<td>What is 2+2?</td>"
|
||||
"<td>4</td>"
|
||||
"<td>correct</td>"
|
||||
"<td>3</td>"
|
||||
"<td>incorrect</td>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"</tr><tr>" # -----
|
||||
"<td>MA</td>"
|
||||
"<td>What C datatypes are 8 bits? (assume i386)</td>"
|
||||
"<td>int</td>"
|
||||
"<td/>"
|
||||
"<td>float</td>"
|
||||
"<td/>"
|
||||
"<td>double</td>"
|
||||
"<td/>"
|
||||
"<td>char</td>"
|
||||
"</tr><tr>" # -----
|
||||
"<td>TF</td>"
|
||||
"<td>Bagpipes are awesome.</td>"
|
||||
"<td>true</td>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"</tr><tr>" # -----
|
||||
"<td>ESS</td>"
|
||||
"<td>How have the original Henry Hornbostel buildings influenced campus architecture and"
|
||||
" design in the last 30 years?</td>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"</tr><tr>" # -----
|
||||
"<td>ORD</td>"
|
||||
"<td>Rank the following in their order of operation.</td>"
|
||||
"<td>Parentheses</td>"
|
||||
"<td>Exponents</td>"
|
||||
"<td>Division</td>"
|
||||
"<td>Addition</td>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"</tr><tr>" # -----
|
||||
"<td>FIB</td>"
|
||||
"<td>The student activities fee is</td>"
|
||||
"<td>95</td>"
|
||||
"<td>dollars for students enrolled in</td>"
|
||||
"<td>19</td>"
|
||||
"<td>units or more,</td>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"<td/>"
|
||||
"</tr><tr>" # -----
|
||||
"<td>MAT</td>"
|
||||
"<td>Match the lower-case greek letter with its capital form.</td>"
|
||||
"<td>λ</td>"
|
||||
"<td>Λ</td>"
|
||||
"<td>α</td>"
|
||||
"<td>γ</td>"
|
||||
"<td>Γ</td>"
|
||||
"<td>φ</td>"
|
||||
"<td>Φ</td>"
|
||||
"</tr></table>"
|
||||
)
|
||||
@@ -0,0 +1,329 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.partition.test_constants import (
|
||||
EXPECTED_TABLE,
|
||||
EXPECTED_TABLE_SEMICOLON_DELIMITER,
|
||||
EXPECTED_TABLE_WITH_EMOJI,
|
||||
EXPECTED_TABLE_WITH_LINE_DELIMITER,
|
||||
EXPECTED_TEXT,
|
||||
EXPECTED_TEXT_SEMICOLON_DELIMITER,
|
||||
EXPECTED_TEXT_WITH_EMOJI,
|
||||
EXPECTED_TEXT_WITH_LINE_DELIMITER,
|
||||
EXPECTED_TEXT_XLSX,
|
||||
)
|
||||
from test_unstructured.unit_utils import (
|
||||
FixtureRequest,
|
||||
Mock,
|
||||
assert_round_trips_through_JSON,
|
||||
example_doc_path,
|
||||
function_mock,
|
||||
)
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.cleaners.core import clean_extra_whitespace
|
||||
from unstructured.documents.elements import Table
|
||||
from unstructured.partition.csv import _CsvPartitioningContext, partition_csv
|
||||
from unstructured.partition.utils.constants import UNSTRUCTURED_INCLUDE_DEBUG_METADATA
|
||||
|
||||
EXPECTED_FILETYPE = "text/csv"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "expected_text", "expected_table"),
|
||||
[
|
||||
("stanley-cups.csv", EXPECTED_TEXT, EXPECTED_TABLE),
|
||||
("stanley-cups-with-emoji.csv", EXPECTED_TEXT_WITH_EMOJI, EXPECTED_TABLE_WITH_EMOJI),
|
||||
(
|
||||
"table-semicolon-delimiter.csv",
|
||||
EXPECTED_TEXT_SEMICOLON_DELIMITER,
|
||||
EXPECTED_TABLE_SEMICOLON_DELIMITER,
|
||||
),
|
||||
(
|
||||
"csv-with-line-delimiter.csv",
|
||||
EXPECTED_TEXT_WITH_LINE_DELIMITER,
|
||||
EXPECTED_TABLE_WITH_LINE_DELIMITER,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_partition_csv_from_filename(filename: str, expected_text: str, expected_table: str):
|
||||
f_path = f"example-docs/{filename}"
|
||||
elements = partition_csv(filename=f_path)
|
||||
|
||||
assert clean_extra_whitespace(elements[0].text) == expected_text
|
||||
assert elements[0].metadata.text_as_html == expected_table
|
||||
assert elements[0].metadata.filetype == EXPECTED_FILETYPE
|
||||
assert elements[0].metadata.filename == filename
|
||||
|
||||
|
||||
@pytest.mark.parametrize("infer_table_structure", [True, False])
|
||||
def test_partition_csv_from_filename_infer_table_structure(infer_table_structure: bool):
|
||||
f_path = "example-docs/stanley-cups.csv"
|
||||
elements = partition_csv(filename=f_path, infer_table_structure=infer_table_structure)
|
||||
|
||||
table_element_has_text_as_html_field = (
|
||||
hasattr(elements[0].metadata, "text_as_html")
|
||||
and elements[0].metadata.text_as_html is not None
|
||||
)
|
||||
assert table_element_has_text_as_html_field == infer_table_structure
|
||||
|
||||
|
||||
def test_partition_csv_from_filename_with_metadata_filename():
|
||||
elements = partition_csv(example_doc_path("stanley-cups.csv"), metadata_filename="test")
|
||||
|
||||
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TEXT
|
||||
assert elements[0].metadata.filename == "test"
|
||||
|
||||
|
||||
def test_partition_csv_with_encoding():
|
||||
elements = partition_csv(example_doc_path("stanley-cups-utf-16.csv"), encoding="utf-16")
|
||||
|
||||
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TEXT
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "expected_text", "expected_table"),
|
||||
[
|
||||
("stanley-cups.csv", EXPECTED_TEXT, EXPECTED_TABLE),
|
||||
("stanley-cups-with-emoji.csv", EXPECTED_TEXT_WITH_EMOJI, EXPECTED_TABLE_WITH_EMOJI),
|
||||
],
|
||||
)
|
||||
def test_partition_csv_from_file(filename: str, expected_text: str, expected_table: str):
|
||||
f_path = f"example-docs/{filename}"
|
||||
with open(f_path, "rb") as f:
|
||||
elements = partition_csv(file=f)
|
||||
assert clean_extra_whitespace(elements[0].text) == expected_text
|
||||
assert isinstance(elements[0], Table)
|
||||
assert elements[0].metadata.text_as_html == expected_table
|
||||
assert elements[0].metadata.filetype == EXPECTED_FILETYPE
|
||||
assert elements[0].metadata.filename is None
|
||||
if UNSTRUCTURED_INCLUDE_DEBUG_METADATA:
|
||||
assert {element.metadata.detection_origin for element in elements} == {"csv"}
|
||||
|
||||
|
||||
def test_partition_csv_from_file_with_metadata_filename():
|
||||
with open(example_doc_path("stanley-cups.csv"), "rb") as f:
|
||||
elements = partition_csv(file=f, metadata_filename="test")
|
||||
|
||||
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TEXT
|
||||
assert elements[0].metadata.filename == "test"
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_csv_from_file_path_gets_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.csv.get_last_modified_date",
|
||||
return_value=filesystem_last_modified,
|
||||
)
|
||||
|
||||
elements = partition_csv(example_doc_path("stanley-cups.csv"))
|
||||
|
||||
assert elements[0].metadata.last_modified == filesystem_last_modified
|
||||
|
||||
|
||||
def test_partition_csv_from_file_path_prefers_metadata_last_modified(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
|
||||
mocker.patch(
|
||||
"unstructured.partition.csv.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_csv(
|
||||
example_doc_path("stanley-cups.csv"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert elements[0].metadata.last_modified == metadata_last_modified
|
||||
|
||||
|
||||
def test_partition_csv_from_file_gets_last_modified_None():
|
||||
with open(example_doc_path("stanley-cups.csv"), "rb") as f:
|
||||
elements = partition_csv(file=f)
|
||||
|
||||
assert elements[0].metadata.last_modified is None
|
||||
|
||||
|
||||
def test_partition_csv_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
|
||||
with open(example_doc_path("stanley-cups.csv"), "rb") as f:
|
||||
elements = partition_csv(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert elements[0].metadata.last_modified == metadata_last_modified
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", ["stanley-cups.csv", "stanley-cups-with-emoji.csv"])
|
||||
def test_partition_csv_with_json(filename: str):
|
||||
elements = partition_csv(filename=example_doc_path(filename))
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
def test_add_chunking_strategy_to_partition_csv_non_default():
|
||||
filename = "example-docs/stanley-cups.csv"
|
||||
|
||||
elements = partition_csv(filename=filename)
|
||||
chunk_elements = partition_csv(
|
||||
filename,
|
||||
chunking_strategy="by_title",
|
||||
max_characters=9,
|
||||
combine_text_under_n_chars=0,
|
||||
include_header=False,
|
||||
)
|
||||
chunks = chunk_by_title(elements, max_characters=9, combine_text_under_n_chars=0)
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
# NOTE (jennings) partition_csv returns a single TableElement per sheet,
|
||||
# so leaving off additional tests for multiple languages like the other partitions
|
||||
def test_partition_csv_element_metadata_has_languages():
|
||||
filename = "example-docs/stanley-cups.csv"
|
||||
elements = partition_csv(filename=filename, strategy="fast", include_header=False)
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_csv_respects_languages_arg():
|
||||
filename = "example-docs/stanley-cups.csv"
|
||||
elements = partition_csv(
|
||||
filename=filename, strategy="fast", languages=["deu"], include_header=False
|
||||
)
|
||||
assert elements[0].metadata.languages == ["deu"]
|
||||
|
||||
|
||||
def test_partition_csv_header():
|
||||
elements = partition_csv(
|
||||
example_doc_path("stanley-cups.csv"), strategy="fast", include_header=True
|
||||
)
|
||||
|
||||
table = elements[0]
|
||||
assert table.text == "Stanley Cups Unnamed: 1 Unnamed: 2 " + EXPECTED_TEXT_XLSX
|
||||
assert table.metadata.text_as_html is not None
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# UNIT-TESTS
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class Describe_CsvPartitioningContext:
|
||||
"""Unit-test suite for `unstructured.partition.csv._CsvPartitioningContext`."""
|
||||
|
||||
# -- .load() ------------------------------------------------
|
||||
|
||||
def it_provides_a_validating_alternate_constructor(self):
|
||||
ctx = _CsvPartitioningContext.load(
|
||||
file_path=example_doc_path("stanley-cups.csv"),
|
||||
file=None,
|
||||
encoding=None,
|
||||
include_header=True,
|
||||
infer_table_structure=True,
|
||||
)
|
||||
assert isinstance(ctx, _CsvPartitioningContext)
|
||||
|
||||
def and_the_validating_constructor_raises_on_an_invalid_context(self):
|
||||
with pytest.raises(ValueError, match="either file-path or file-like object must be prov"):
|
||||
_CsvPartitioningContext.load(
|
||||
file_path=None,
|
||||
file=None,
|
||||
encoding=None,
|
||||
include_header=True,
|
||||
infer_table_structure=True,
|
||||
)
|
||||
|
||||
# -- .delimiter ---------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file_name",
|
||||
[
|
||||
"stanley-cups.csv",
|
||||
# -- Issue #2643: previously raised `_csv.Error: Could not determine delimiter` on
|
||||
# -- this file
|
||||
"csv-with-long-lines.csv",
|
||||
],
|
||||
)
|
||||
def it_auto_detects_the_delimiter_for_a_comma_delimited_CSV_file(self, file_name: str):
|
||||
ctx = _CsvPartitioningContext(example_doc_path(file_name))
|
||||
assert ctx.delimiter == ","
|
||||
|
||||
def and_it_auto_detects_the_delimiter_for_a_semicolon_delimited_CSV_file(self):
|
||||
ctx = _CsvPartitioningContext(example_doc_path("semicolon-delimited.csv"))
|
||||
assert ctx.delimiter == ";"
|
||||
|
||||
def but_it_returns_None_as_the_delimiter_for_a_single_column_CSV_file(self):
|
||||
ctx = _CsvPartitioningContext(example_doc_path("single-column.csv"))
|
||||
assert ctx.delimiter is None
|
||||
|
||||
# -- .header ------------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(("include_header", "expected_value"), [(False, None), (True, 0)])
|
||||
def it_identifies_the_header_row_based_on_include_header_arg(
|
||||
self, include_header: bool, expected_value: int | None
|
||||
):
|
||||
assert _CsvPartitioningContext(include_header=include_header).header == expected_value
|
||||
|
||||
# -- .last_modified -----------------------------------------
|
||||
|
||||
def it_gets_last_modified_from_the_filesystem_when_a_path_is_provided(
|
||||
self, get_last_modified_date_: Mock
|
||||
):
|
||||
filesystem_last_modified = "2024-08-04T02:23:53"
|
||||
get_last_modified_date_.return_value = filesystem_last_modified
|
||||
ctx = _CsvPartitioningContext(file_path="a/b/document.csv")
|
||||
|
||||
last_modified = ctx.last_modified
|
||||
|
||||
get_last_modified_date_.assert_called_once_with("a/b/document.csv")
|
||||
assert last_modified == filesystem_last_modified
|
||||
|
||||
def and_it_falls_back_to_None_for_the_last_modified_date_when_file_path_is_not_provided(self):
|
||||
file = io.BytesIO(b"abcdefg")
|
||||
ctx = _CsvPartitioningContext(file=file)
|
||||
|
||||
last_modified = ctx.last_modified
|
||||
|
||||
assert last_modified is None
|
||||
|
||||
# -- .open() ------------------------------------------------
|
||||
|
||||
def it_provides_transparent_access_to_the_source_file_when_it_is_a_file_like_object(self):
|
||||
with open(example_doc_path("stanley-cups.csv"), "rb") as f:
|
||||
# -- read so file cursor is at end of file --
|
||||
f.read()
|
||||
ctx = _CsvPartitioningContext(file=f)
|
||||
with ctx.open() as file:
|
||||
assert file is f
|
||||
# -- read cursor is reset to 0 on .open() context entry --
|
||||
assert f.tell() == 0
|
||||
assert file.read(14) == b"Stanley Cups,,"
|
||||
assert f.tell() == 14
|
||||
|
||||
# -- and read cursor is reset to 0 on .open() context exit --
|
||||
assert f.tell() == 0
|
||||
|
||||
def it_provides_transparent_access_to_the_source_file_when_it_is_a_file_path(self):
|
||||
ctx = _CsvPartitioningContext(example_doc_path("stanley-cups.csv"))
|
||||
with ctx.open() as file:
|
||||
assert file.read(14) == b"Stanley Cups,,"
|
||||
|
||||
# -- .validate() --------------------------------------------
|
||||
|
||||
def it_raises_when_neither_file_path_nor_file_is_provided(self):
|
||||
with pytest.raises(ValueError, match="either file-path or file-like object must be prov"):
|
||||
_CsvPartitioningContext()._validate()
|
||||
|
||||
# -- fixtures --------------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture()
|
||||
def get_last_modified_date_(self, request: FixtureRequest) -> Mock:
|
||||
return function_mock(request, "unstructured.partition.csv.get_last_modified_date")
|
||||
@@ -0,0 +1,283 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
"""Test suite for `unstructured.partition.doc` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pathlib
|
||||
from typing import Any, Iterator
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.unit_utils import (
|
||||
ANY,
|
||||
CaptureFixture,
|
||||
FixtureRequest,
|
||||
assert_round_trips_through_JSON,
|
||||
example_doc_path,
|
||||
method_mock,
|
||||
)
|
||||
from unstructured.chunking.basic import chunk_elements
|
||||
from unstructured.documents.elements import (
|
||||
Address,
|
||||
CompositeElement,
|
||||
Element,
|
||||
ListItem,
|
||||
NarrativeText,
|
||||
Table,
|
||||
TableChunk,
|
||||
Text,
|
||||
Title,
|
||||
)
|
||||
from unstructured.partition.doc import partition_doc
|
||||
from unstructured.partition.docx import partition_docx
|
||||
|
||||
|
||||
def test_partition_doc_matches_partition_docx(request: FixtureRequest):
|
||||
doc_file_path = example_doc_path("simple.doc")
|
||||
docx_file_path = example_doc_path("simple.docx")
|
||||
|
||||
assert partition_doc(doc_file_path) == partition_docx(docx_file_path)
|
||||
|
||||
|
||||
# -- document-source (file or filename) ----------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_doc_from_filename(expected_elements: list[Element], capsys: CaptureFixture[str]):
|
||||
elements = partition_doc(example_doc_path("simple.doc"))
|
||||
|
||||
assert elements == expected_elements
|
||||
assert all(e.metadata.file_directory == example_doc_path("") for e in elements)
|
||||
assert capsys.readouterr().out == ""
|
||||
assert capsys.readouterr().err == ""
|
||||
|
||||
|
||||
def test_partition_doc_from_file_with_libre_office_filter(
|
||||
expected_elements: list[Element], capsys: CaptureFixture[str]
|
||||
):
|
||||
with open(example_doc_path("simple.doc"), "rb") as f:
|
||||
elements = partition_doc(file=f, libre_office_filter="MS Word 2007 XML")
|
||||
|
||||
assert elements == expected_elements
|
||||
assert capsys.readouterr().out == ""
|
||||
assert capsys.readouterr().err == ""
|
||||
|
||||
|
||||
def test_partition_doc_from_file_with_no_libre_office_filter(
|
||||
expected_elements: list[Element], capsys: CaptureFixture[str]
|
||||
):
|
||||
with open(example_doc_path("simple.doc"), "rb") as f:
|
||||
elements = partition_doc(file=f, libre_office_filter=None)
|
||||
|
||||
assert elements == expected_elements
|
||||
assert capsys.readouterr().out == ""
|
||||
assert capsys.readouterr().err == ""
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_doc_raises_when_both_a_filename_and_file_are_specified():
|
||||
doc_file_path = example_doc_path("simple.doc")
|
||||
|
||||
with open(doc_file_path, "rb") as f:
|
||||
with pytest.raises(ValueError, match="Exactly one of filename and file must be specified"):
|
||||
partition_doc(filename=doc_file_path, file=f)
|
||||
|
||||
|
||||
def test_partition_doc_raises_when_neither_a_file_path_nor_a_file_like_object_are_provided():
|
||||
with pytest.raises(ValueError, match="Exactly one of filename and file must be specified"):
|
||||
partition_doc()
|
||||
|
||||
|
||||
def test_partition_raises_with_missing_doc(tmp_path: pathlib.Path):
|
||||
doc_filename = str(tmp_path / "asdf.doc")
|
||||
|
||||
with pytest.raises(ValueError, match="asdf.doc does not exist"):
|
||||
partition_doc(filename=doc_filename)
|
||||
|
||||
|
||||
# -- .metadata.filename --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_doc_from_filename_gets_filename_from_filename_arg():
|
||||
elements = partition_doc(example_doc_path("simple.doc"))
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename == "simple.doc" for e in elements)
|
||||
|
||||
|
||||
def test_partition_doc_from_file_gets_filename_None():
|
||||
with open(example_doc_path("simple.doc"), "rb") as f:
|
||||
elements = partition_doc(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_doc_from_filename_prefers_metadata_filename():
|
||||
elements = partition_doc(example_doc_path("simple.doc"), metadata_filename="test")
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(element.metadata.filename == "test" for element in elements)
|
||||
|
||||
|
||||
def test_partition_doc_from_file_prefers_metadata_filename():
|
||||
with open(example_doc_path("simple.doc"), "rb") as f:
|
||||
elements = partition_doc(file=f, metadata_filename="test")
|
||||
|
||||
assert all(e.metadata.filename == "test" for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.filetype --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_doc_gets_the_DOC_MIME_type_in_metadata_filetype():
|
||||
DOC_MIME_TYPE = "application/msword"
|
||||
elements = partition_doc(example_doc_path("simple.doc"))
|
||||
assert all(e.metadata.filetype == DOC_MIME_TYPE for e in elements), (
|
||||
f"Expected all elements to have '{DOC_MIME_TYPE}' as their filetype, but got:"
|
||||
f" {repr(elements[0].metadata.filetype)}"
|
||||
)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_doc_pulls_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.doc.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_doc(example_doc_path("fake.doc"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_doc_prefers_metadata_last_modified_when_provided(
|
||||
mocker: MockFixture,
|
||||
):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.doc.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_doc(
|
||||
example_doc_path("simple.doc"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# -- language-recognition metadata ---------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_doc_adds_languages_metadata():
|
||||
elements = partition_doc(example_doc_path("simple.doc"))
|
||||
assert all(e.metadata.languages == ["eng"] for e in elements)
|
||||
|
||||
|
||||
def test_partition_doc_respects_detect_language_per_element_arg():
|
||||
elements = partition_doc(
|
||||
example_doc_path("language-docs/eng_spa_mult.doc"), detect_language_per_element=True
|
||||
)
|
||||
assert [e.metadata.languages for e in elements] == [
|
||||
["eng"],
|
||||
["spa", "eng"],
|
||||
["eng"],
|
||||
["eng"],
|
||||
["spa"],
|
||||
]
|
||||
|
||||
|
||||
# -- miscellaneous -------------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "expected_value"),
|
||||
[({}, "hi_res"), ({"strategy": None}, "hi_res"), ({"strategy": "auto"}, "auto")],
|
||||
)
|
||||
def test_partition_odt_forwards_strategy_arg_to_partition_docx(
|
||||
request: FixtureRequest, kwargs: dict[str, Any], expected_value: str | None
|
||||
):
|
||||
from unstructured.partition.docx import _DocxPartitioner
|
||||
|
||||
def fake_iter_document_elements(self: _DocxPartitioner) -> Iterator[Element]:
|
||||
yield Text(f"strategy == {self._opts.strategy}")
|
||||
|
||||
_iter_elements_ = method_mock(
|
||||
request,
|
||||
_DocxPartitioner,
|
||||
"_iter_document_elements",
|
||||
side_effect=fake_iter_document_elements,
|
||||
)
|
||||
|
||||
(element,) = partition_doc(example_doc_path("simple.doc"), **kwargs)
|
||||
|
||||
_iter_elements_.assert_called_once_with(ANY)
|
||||
assert element.text == f"strategy == {expected_value}"
|
||||
|
||||
|
||||
def test_partition_doc_grabs_emphasized_texts():
|
||||
expected_emphasized_text_contents = ["bold", "italic", "bold-italic", "bold-italic"]
|
||||
expected_emphasized_text_tags = ["b", "i", "b", "i"]
|
||||
|
||||
elements = partition_doc(example_doc_path("fake-doc-emphasized-text.doc"))
|
||||
|
||||
assert isinstance(elements[0], Table)
|
||||
assert elements[0].metadata.emphasized_text_contents == expected_emphasized_text_contents
|
||||
assert elements[0].metadata.emphasized_text_tags == expected_emphasized_text_tags
|
||||
|
||||
assert elements[1] == NarrativeText("I am a bold italic bold-italic text.")
|
||||
assert elements[1].metadata.emphasized_text_contents == expected_emphasized_text_contents
|
||||
assert elements[1].metadata.emphasized_text_tags == expected_emphasized_text_tags
|
||||
|
||||
assert elements[2] == NarrativeText("I am a normal text.")
|
||||
assert elements[2].metadata.emphasized_text_contents is None
|
||||
assert elements[2].metadata.emphasized_text_tags is None
|
||||
|
||||
|
||||
def test_partition_doc_round_trips_through_json():
|
||||
"""Elements produced can be serialized then deserialized without loss."""
|
||||
assert_round_trips_through_JSON(partition_doc(example_doc_path("simple.doc")))
|
||||
|
||||
|
||||
def test_partition_doc_chunks_elements_when_chunking_strategy_is_specified():
|
||||
document_path = example_doc_path("simple.doc")
|
||||
elements = partition_doc(document_path)
|
||||
chunks = partition_doc(document_path, chunking_strategy="basic")
|
||||
|
||||
# -- all chunks are chunk element-types --
|
||||
assert all(isinstance(c, (CompositeElement, Table, TableChunk)) for c in chunks)
|
||||
# -- chunks from partitioning match those produced by chunking elements in separate step --
|
||||
assert chunks == chunk_elements(elements)
|
||||
|
||||
|
||||
def test_partition_doc_assigns_deterministic_and_unique_element_ids():
|
||||
document_path = example_doc_path("duplicate-paragraphs.doc")
|
||||
|
||||
ids = [element.id for element in partition_doc(document_path)]
|
||||
ids_2 = [element.id for element in partition_doc(document_path)]
|
||||
|
||||
# -- ids should match even though partitioned separately --
|
||||
assert ids == ids_2
|
||||
# -- ids should be unique --
|
||||
assert len(ids) == len(set(ids))
|
||||
|
||||
|
||||
# == module-level fixtures =======================================================================
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def expected_elements() -> list[Element]:
|
||||
return [
|
||||
Title("These are a few of my favorite things:"),
|
||||
ListItem("Parrots"),
|
||||
ListItem("Hockey"),
|
||||
Text("Analysis"),
|
||||
NarrativeText("This is my first thought. This is my second thought."),
|
||||
NarrativeText("This is my third thought."),
|
||||
Text("2023"),
|
||||
Address("DOYLESTOWN, PA 18901"),
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,631 @@
|
||||
"""Test suite for `unstructured.partition.email` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import tempfile
|
||||
from email.message import EmailMessage
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from test_unstructured.unit_utils import (
|
||||
FixtureRequest,
|
||||
Mock,
|
||||
assert_round_trips_through_JSON,
|
||||
example_doc_path,
|
||||
function_mock,
|
||||
)
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import (
|
||||
CompositeElement,
|
||||
ListItem,
|
||||
NarrativeText,
|
||||
Table,
|
||||
TableChunk,
|
||||
Text,
|
||||
Title,
|
||||
)
|
||||
from unstructured.partition.email import EmailPartitioningContext, partition_email
|
||||
|
||||
EXPECTED_OUTPUT = [
|
||||
NarrativeText(text="This is a test email to use for unit tests."),
|
||||
Text(text="Important points:"),
|
||||
ListItem(text="Roses are red"),
|
||||
ListItem(text="Violets are blue"),
|
||||
]
|
||||
|
||||
|
||||
def test_partition_email_from_filename_can_partition_an_RFC_822_email():
|
||||
assert partition_email(example_doc_path("eml/simple-rfc-822.eml")) == [
|
||||
NarrativeText("This is an RFC 822 email message."),
|
||||
NarrativeText(
|
||||
"An RFC 822 message is characterized by its simple, text-based format, which includes"
|
||||
' a header and a body. The header contains structured fields such as "From", "To",'
|
||||
' "Date", and "Subject", each followed by a colon and the corresponding information.'
|
||||
" The body follows the header, separated by a blank line, and contains the main"
|
||||
" content of the email."
|
||||
),
|
||||
NarrativeText(
|
||||
"The structure ensures compatibility and readability across different email systems"
|
||||
" and clients, adhering to the standards set by the Internet Engineering Task Force"
|
||||
" (IETF)."
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_partition_email_from_file_can_partition_an_email():
|
||||
with open(example_doc_path("eml/fake-email.eml"), "rb") as f:
|
||||
assert partition_email(file=f) == EXPECTED_OUTPUT
|
||||
|
||||
|
||||
def test_partition_email_from_spooled_temp_file_can_partition_an_email():
|
||||
with tempfile.SpooledTemporaryFile() as file:
|
||||
with open(example_doc_path("eml/fake-email.eml"), "rb") as f:
|
||||
file.write(f.read())
|
||||
file.seek(0)
|
||||
|
||||
assert partition_email(file=file) == EXPECTED_OUTPUT
|
||||
|
||||
|
||||
def test_partition_email_can_partition_an_HTML_only_email_with_Base64_ISO_8859_1_charset():
|
||||
assert partition_email(example_doc_path("eml/mime-html-only.eml")) == [
|
||||
NarrativeText("This is a text/html part."),
|
||||
NarrativeText(
|
||||
"The first emoticon, :) , was proposed by Scott Fahlman in 1982 to indicate just or"
|
||||
" sarcasm in text emails."
|
||||
),
|
||||
NarrativeText(
|
||||
"Gmail was launched by Google in 2004 with 1 GB of free storage, significantly more"
|
||||
" than what other services offered at the time."
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_extract_email_from_text_plain_matches_elements_extracted_from_text_html():
|
||||
file_path = example_doc_path("eml/fake-email.eml")
|
||||
|
||||
elements_from_text = partition_email(file_path, content_source="text/plain")
|
||||
elements_from_html = partition_email(file_path, content_source="text/html")
|
||||
|
||||
assert all(e.text == eo.text for e, eo in zip(elements_from_text, EXPECTED_OUTPUT))
|
||||
assert elements_from_html == EXPECTED_OUTPUT
|
||||
assert all(eh.text == et.text for eh, et in zip(elements_from_html, elements_from_text))
|
||||
|
||||
|
||||
def test_partition_email_round_trips_via_json():
|
||||
elements = partition_email(example_doc_path("eml/fake-email.eml"))
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
# -- transfer-encodings --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_email_partitions_an_HTML_part_with_Base64_encoded_UTF_8_charset():
|
||||
assert partition_email(example_doc_path("eml/fake-email-b64.eml")) == EXPECTED_OUTPUT
|
||||
|
||||
|
||||
def test_partition_email_partitions_a_text_plain_part_with_Base64_encoded_windows_1255_charset():
|
||||
elements = partition_email(
|
||||
example_doc_path("eml/email-no-utf8-2008-07-16.062410.eml"),
|
||||
content_source="text/plain",
|
||||
)
|
||||
|
||||
assert len(elements) == 30
|
||||
assert elements[1].text.startswith("אני חושב שזה לא יהיה מקצועי והוגן שאני אראה לך היכן")
|
||||
|
||||
|
||||
def test_partition_email_partitions_an_html_part_with_quoted_printable_encoded_ISO_8859_1_charset():
|
||||
elements = partition_email(
|
||||
example_doc_path("eml/email-no-utf8-2014-03-17.111517.eml"),
|
||||
content_source="text/html",
|
||||
process_attachments=False,
|
||||
)
|
||||
|
||||
assert len(elements) == 1
|
||||
assert isinstance(elements[0], Table)
|
||||
assert elements[0].text.startswith("Slava Gxyzxyz Hi Slava, The password for your Google")
|
||||
|
||||
|
||||
# -- edge-cases ----------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_email_accepts_a_whitespace_only_file():
|
||||
"""Should produce no elements but should not raise an exception."""
|
||||
assert partition_email(example_doc_path("eml/empty.eml")) == []
|
||||
|
||||
|
||||
def test_partition_email_can_partition_an_empty_email():
|
||||
assert (
|
||||
partition_email(example_doc_path("eml/mime-no-body.eml"), process_attachments=False) == []
|
||||
)
|
||||
|
||||
|
||||
def test_partition_email_does_not_break_on_an_encrypted_message():
|
||||
assert (
|
||||
partition_email(example_doc_path("eml/fake-encrypted.eml"), process_attachments=False) == []
|
||||
)
|
||||
|
||||
|
||||
def test_partition_email_finds_content_when_it_is_marked_with_content_disposition_inline():
|
||||
elements = partition_email(
|
||||
example_doc_path("eml/email-inline-content-disposition.eml"), process_attachments=False
|
||||
)
|
||||
|
||||
assert len(elements) == 1
|
||||
e = elements[0]
|
||||
assert isinstance(e, Text)
|
||||
assert e.text == "This is a test of inline"
|
||||
|
||||
|
||||
def test_partition_email_from_filename_malformed_encoding():
|
||||
elements = partition_email(filename=example_doc_path("eml/fake-email-malformed-encoding.eml"))
|
||||
assert elements == EXPECTED_OUTPUT
|
||||
|
||||
|
||||
# -- error behaviors -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_email_raises_when_no_message_source_is_specified():
|
||||
with pytest.raises(ValueError, match="no document specified; either a `filename` or `file`"):
|
||||
partition_email()
|
||||
|
||||
|
||||
def test_partition_email_raises_with_invalid_content_type():
|
||||
with pytest.raises(ValueError, match="'application/json' is not a valid value for content_s"):
|
||||
partition_email(example_doc_path("eml/fake-email.eml"), content_source="application/json")
|
||||
|
||||
|
||||
# -- .metadata -----------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_email_augments_message_body_elements_with_email_metadata():
|
||||
elements = partition_email(example_doc_path("eml/mime-multi-to-cc-bcc.eml"))
|
||||
|
||||
assert all(
|
||||
e.metadata.bcc_recipient == ["John <john@example.com>", "Mary <mary@example.com>"]
|
||||
for e in elements
|
||||
)
|
||||
assert all(
|
||||
e.metadata.cc_recipient == ["Tom <tom@example.com>", "Alice <alice@example.com>"]
|
||||
for e in elements
|
||||
)
|
||||
assert all(e.metadata.email_message_id == "2143658709@example.com" for e in elements)
|
||||
assert all(e.metadata.sent_from == ["sender@example.com"] for e in elements)
|
||||
assert all(
|
||||
e.metadata.sent_to == ["Bob <bob@example.com>", "Sue <sue@example.com>"] for e in elements
|
||||
)
|
||||
assert all(e.metadata.subject == "Example Plain-Text MIME Message" for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.filename --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_email_from_filename_gets_filename_metadata_from_file_path():
|
||||
elements = partition_email(example_doc_path("eml/fake-email.eml"))
|
||||
|
||||
assert all(e.metadata.filename == "fake-email.eml" for e in elements)
|
||||
assert all(e.metadata.file_directory == example_doc_path("eml") for e in elements)
|
||||
|
||||
|
||||
def test_partition_email_from_file_gets_filename_metadata_None():
|
||||
with open(example_doc_path("eml/fake-email.eml"), "rb") as f:
|
||||
elements = partition_email(file=f)
|
||||
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
assert all(e.metadata.file_directory is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_email_from_filename_prefers_metadata_filename():
|
||||
elements = partition_email(
|
||||
example_doc_path("eml/fake-email.eml"), metadata_filename="a/b/c.eml"
|
||||
)
|
||||
|
||||
assert all(e.metadata.filename == "c.eml" for e in elements)
|
||||
assert all(e.metadata.file_directory == "a/b" for e in elements)
|
||||
|
||||
|
||||
def test_partition_email_from_file_prefers_metadata_filename():
|
||||
with open(example_doc_path("eml/fake-email.eml"), "rb") as f:
|
||||
elements = partition_email(file=f, metadata_filename="d/e/f.eml")
|
||||
|
||||
assert all(e.metadata.filename == "f.eml" for e in elements)
|
||||
assert all(e.metadata.file_directory == "d/e" for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.filetype --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_email_gets_the_EML_MIME_type_in_metadata_filetype_for_message_body_elements():
|
||||
EML_MIME_TYPE = "message/rfc822"
|
||||
elements = partition_email(example_doc_path("eml/fake-email.eml"))
|
||||
assert all(e.metadata.filetype == EML_MIME_TYPE for e in elements), (
|
||||
f"Expected all elements to have '{EML_MIME_TYPE}' as their filetype, but got:"
|
||||
f" {repr(elements[0].metadata.filetype)}"
|
||||
)
|
||||
|
||||
|
||||
# -- .metadata.languages -------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_email_element_metadata_has_languages():
|
||||
elements = partition_email(example_doc_path("eml/fake-email.eml"))
|
||||
assert all(e.metadata.languages == ["eng"] for e in elements)
|
||||
|
||||
|
||||
def test_partition_email_respects_languages_arg():
|
||||
elements = partition_email(example_doc_path("eml/fake-email.eml"), languages=["deu"])
|
||||
assert all(element.metadata.languages == ["deu"] for element in elements)
|
||||
|
||||
|
||||
def test_partition_eml_respects_detect_language_per_element():
|
||||
elements = partition_email(
|
||||
example_doc_path("language-docs/eng_spa_mult.eml"),
|
||||
detect_language_per_element=True,
|
||||
)
|
||||
# languages other than English and Spanish are detected by this partitioner,
|
||||
# so this test is slightly different from the other partition tests
|
||||
langs = {e.metadata.languages[0] for e in elements if e.metadata.languages is not None}
|
||||
|
||||
assert "eng" in langs
|
||||
assert "spa" in langs
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_email_from_file_path_gets_last_modified_from_Date_header():
|
||||
elements = partition_email(example_doc_path("eml/fake-email.eml"))
|
||||
assert all(e.metadata.last_modified == "2022-12-16T22:04:16+00:00" for e in elements)
|
||||
|
||||
|
||||
def test_partition_email_from_file_gets_last_modified_from_Date_header():
|
||||
with open(example_doc_path("eml/fake-email.eml"), "rb") as f:
|
||||
elements = partition_email(file=f)
|
||||
|
||||
assert all(e.metadata.last_modified == "2022-12-16T22:04:16+00:00" for e in elements)
|
||||
|
||||
|
||||
def test_partition_email_from_file_path_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
|
||||
elements = partition_email(
|
||||
example_doc_path("eml/fake-email.eml"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_email_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
with open(example_doc_path("eml/fake-email.eml"), "rb") as f:
|
||||
elements = partition_email(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# -- chunking ------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_email_chunks_when_so_instructed():
|
||||
"""Note it's actually the delegate partitioners that do the chunking."""
|
||||
elements = partition_email(example_doc_path("eml/fake-email.txt"))
|
||||
chunks = partition_email(example_doc_path("eml/fake-email.txt"), chunking_strategy="by_title")
|
||||
separately_chunked_chunks = chunk_by_title(elements)
|
||||
|
||||
assert all(isinstance(c, (CompositeElement, Table, TableChunk)) for c in chunks)
|
||||
assert chunks != elements
|
||||
assert chunks == separately_chunked_chunks
|
||||
|
||||
|
||||
def test_partition_email_chunks_attachments_too():
|
||||
chunks = partition_email(
|
||||
example_doc_path("eml/fake-email-attachment.eml"),
|
||||
chunking_strategy="by_title",
|
||||
process_attachments=True,
|
||||
)
|
||||
|
||||
assert len(chunks) == 2
|
||||
assert all(isinstance(c, CompositeElement) for c in chunks)
|
||||
attachment_chunk = chunks[-1]
|
||||
assert attachment_chunk.text == "Hey this is a fake attachment!"
|
||||
assert attachment_chunk.metadata.filename == "fake-attachment.txt"
|
||||
assert attachment_chunk.metadata.attached_to_filename == "fake-email-attachment.eml"
|
||||
assert all(c.metadata.last_modified == "2022-12-23T18:08:48+00:00" for c in chunks)
|
||||
|
||||
|
||||
# -- attachments ---------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_email_also_partitions_attachments_when_so_instructed():
|
||||
elements = partition_email(
|
||||
example_doc_path("eml/email-equals-attachment-filename.eml"), process_attachments=True
|
||||
)
|
||||
|
||||
assert elements == [
|
||||
NarrativeText("Below is an example of an odd filename"),
|
||||
Title("Odd filename"),
|
||||
]
|
||||
|
||||
|
||||
def test_partition_email_can_process_attachments():
|
||||
elements = partition_email(
|
||||
example_doc_path("eml/fake-email-attachment.eml"), process_attachments=True
|
||||
)
|
||||
|
||||
assert elements == [
|
||||
Text("Hello!"),
|
||||
NarrativeText("Here's the attachments!"),
|
||||
NarrativeText("It includes:"),
|
||||
ListItem("Lots of whitespace"),
|
||||
ListItem("Little to no content"),
|
||||
ListItem("and is a quick read"),
|
||||
Text("Best,"),
|
||||
Text("Mallori"),
|
||||
NarrativeText("Hey this is a fake attachment!"),
|
||||
]
|
||||
assert all(e.metadata.last_modified == "2022-12-23T18:08:48+00:00" for e in elements)
|
||||
attachment_element = elements[-1]
|
||||
assert attachment_element.text == "Hey this is a fake attachment!"
|
||||
assert attachment_element.metadata.filename == "fake-attachment.txt"
|
||||
assert attachment_element.metadata.attached_to_filename == "fake-email-attachment.eml"
|
||||
|
||||
|
||||
def test_partition_email_silently_skips_attachments_it_cannot_partition():
|
||||
elements = partition_email(
|
||||
example_doc_path("eml/mime-attach-mp3.eml"), process_attachments=True
|
||||
)
|
||||
|
||||
# -- no exception is raised --
|
||||
assert elements == [
|
||||
# -- the email body is partitioned --
|
||||
NarrativeText("This is an email with an MP3 attachment."),
|
||||
# -- no elements appear for the attachment --
|
||||
]
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# ISOLATED UNIT TESTS
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class DescribeEmailPartitionerOptions:
|
||||
"""Unit-test suite for `unstructured.partition.email.EmailPartitioningContext` objects."""
|
||||
|
||||
# -- .load() ---------------------------------
|
||||
|
||||
def it_provides_a_validating_constructor(self, ctx_args: dict[str, Any]):
|
||||
ctx_args["file_path"] = example_doc_path("eml/fake-email.eml")
|
||||
|
||||
ctx = EmailPartitioningContext.load(**ctx_args)
|
||||
|
||||
assert isinstance(ctx, EmailPartitioningContext)
|
||||
|
||||
def but_it_raises_when_no_source_document_was_specified(self, ctx_args: dict[str, Any]):
|
||||
with pytest.raises(ValueError, match="no document specified; either a `filename` or `fi"):
|
||||
EmailPartitioningContext.load(**ctx_args)
|
||||
|
||||
def and_it_raises_when_a_file_open_for_reading_str_is_used(self, ctx_args: dict[str, Any]):
|
||||
ctx_args["file"] = io.StringIO("abcdefg")
|
||||
with pytest.raises(ValueError, match="file object must be opened in binary mode"):
|
||||
EmailPartitioningContext.load(**ctx_args)
|
||||
|
||||
def and_it_raises_when_an_invalid_content_source_is_specified(self, ctx_args: dict[str, Any]):
|
||||
ctx_args["file_path"] = example_doc_path("eml/fake-email.eml")
|
||||
ctx_args["content_source"] = "application/json"
|
||||
|
||||
with pytest.raises(ValueError, match="'application/json' is not a valid value for conte"):
|
||||
EmailPartitioningContext.load(**ctx_args)
|
||||
|
||||
# -- .bcc_addresses --------------------------
|
||||
|
||||
def it_provides_access_to_the_Bcc_addresses_when_present(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/mime-multi-to-cc-bcc.eml"))
|
||||
assert ctx.bcc_addresses == ["John <john@example.com>", "Mary <mary@example.com>"]
|
||||
|
||||
def but_it_returns_None_when_there_are_no_Bcc_addresses(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/simple-rfc-822.eml"))
|
||||
assert ctx.bcc_addresses is None
|
||||
|
||||
# -- .body_part ------------------------------
|
||||
|
||||
def it_returns_the_html_body_part_when_there_is_one_by_default(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/mime-different-plain-html.eml"))
|
||||
|
||||
body_part = ctx.body_part
|
||||
|
||||
assert isinstance(body_part, EmailMessage)
|
||||
content = body_part.get_content()
|
||||
assert isinstance(content, str)
|
||||
assert content.startswith("<!DOCTYPE html>")
|
||||
|
||||
def but_it_returns_the_plain_text_body_part_when_there_is_one_when_so_requested(self):
|
||||
ctx = EmailPartitioningContext(
|
||||
example_doc_path("eml/mime-different-plain-html.eml"), content_source="text/plain"
|
||||
)
|
||||
|
||||
body_part = ctx.body_part
|
||||
|
||||
assert isinstance(body_part, EmailMessage)
|
||||
content = body_part.get_content()
|
||||
assert isinstance(content, str)
|
||||
assert content.startswith("This is the text/plain part.")
|
||||
|
||||
def and_it_returns_None_when_the_email_has_no_body(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/mime-no-body.eml"))
|
||||
assert ctx.body_part is None
|
||||
|
||||
# -- .cc_addresses ---------------------------
|
||||
|
||||
def it_provides_access_to_the_Cc_addresses_when_present(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/mime-multi-to-cc-bcc.eml"))
|
||||
assert ctx.cc_addresses == ["Tom <tom@example.com>", "Alice <alice@example.com>"]
|
||||
|
||||
def but_it_returns_None_when_there_are_no_Cc_addresses(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/simple-rfc-822.eml"))
|
||||
assert ctx.cc_addresses is None
|
||||
|
||||
# -- .content_type_preference ----------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("content_source", "expected_value"),
|
||||
[
|
||||
("text/html", ("html", "plain")),
|
||||
("text/plain", ("plain", "html")),
|
||||
],
|
||||
)
|
||||
def it_knows_whether_the_caller_prefers_the_HTML_or_plain_text_body(
|
||||
self, content_source: str, expected_value: tuple[str, ...]
|
||||
):
|
||||
ctx = EmailPartitioningContext(content_source=content_source)
|
||||
assert ctx.content_type_preference == expected_value
|
||||
|
||||
def and_it_defaults_to_preferring_the_HTML_body(self):
|
||||
ctx = EmailPartitioningContext()
|
||||
assert ctx.content_type_preference == ("html", "plain")
|
||||
|
||||
# -- .from -----------------------------------
|
||||
|
||||
def it_knows_the_From_address_of_the_email(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/mime-simple.eml"))
|
||||
assert ctx.from_address == "sender@example.com"
|
||||
|
||||
# -- .message_id -----------------------------
|
||||
|
||||
def it_provides_access_to_the_Message_ID_when_present(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/mime-simple.eml"))
|
||||
assert ctx.message_id == "1234567890@example.com"
|
||||
|
||||
def but_it_returns_None_when_there_is_no_Message_ID_header(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/simple-rfc-822.eml"))
|
||||
assert ctx.message_id is None
|
||||
|
||||
# -- .metadata_file_path ---------------------
|
||||
|
||||
def it_uses_the_metadata_file_path_arg_value_when_one_was_provided(self):
|
||||
ctx = EmailPartitioningContext(metadata_file_path="a/b/c.eml")
|
||||
assert ctx.metadata_file_path == "a/b/c.eml"
|
||||
|
||||
def and_it_uses_the_file_path_arg_value_when_metadata_file_path_was_not_provided(self):
|
||||
ctx = EmailPartitioningContext(file_path="x/y/z.eml")
|
||||
assert ctx.metadata_file_path == "x/y/z.eml"
|
||||
|
||||
def and_it_returns_None_when_neither_file_path_was_provided(self):
|
||||
ctx = EmailPartitioningContext()
|
||||
assert ctx.metadata_file_path is None
|
||||
|
||||
# -- .metadata_last_modified -----------------
|
||||
|
||||
def it_uses_the_metadata_last_modified_arg_value_when_one_was_provided(self):
|
||||
metadata_last_modified = "2023-04-08T12:18:07"
|
||||
ctx = EmailPartitioningContext(metadata_last_modified=metadata_last_modified)
|
||||
assert ctx.metadata_last_modified == metadata_last_modified
|
||||
|
||||
def and_it_uses_the_msg_Date_header_date_when_metadata_last_modified_was_not_provided(
|
||||
self,
|
||||
):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/simple-rfc-822.eml"))
|
||||
assert ctx.metadata_last_modified == "2024-10-01T17:34:56+00:00"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("date_format", "expected_date"),
|
||||
[
|
||||
("test-iso-8601-date.eml", "2025-07-29T12:42:06+00:00"),
|
||||
("test-rfc2822-date.eml", "2025-07-29T12:42:06+00:00"),
|
||||
],
|
||||
)
|
||||
def and_it_correctly_parses_various_date_formats_like_the_ones_that_occur_in_the_wild(
|
||||
self, date_format: str, expected_date: str
|
||||
):
|
||||
ctx = EmailPartitioningContext(example_doc_path(f"eml/{date_format}"))
|
||||
assert ctx.metadata_last_modified == expected_date
|
||||
|
||||
def and_it_returns_none_when_date_header_is_invalid(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/test-invalid-date.eml"))
|
||||
assert ctx._sent_date is None
|
||||
|
||||
def and_it_falls_back_to_filesystem_last_modified_when_no_Date_header_is_present(
|
||||
self, get_last_modified_date_: Mock
|
||||
):
|
||||
"""Not an expected case as according to RFC 5322, the Date header is required."""
|
||||
filesystem_last_modified = "2024-07-09T14:08:17"
|
||||
get_last_modified_date_.return_value = filesystem_last_modified
|
||||
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/rfc822-no-date.eml"))
|
||||
|
||||
assert ctx.metadata_last_modified == filesystem_last_modified
|
||||
|
||||
def and_it_returns_None_when_no_last_modified_is_available(self):
|
||||
with open(example_doc_path("eml/rfc822-no-date.eml"), "rb") as f:
|
||||
ctx = EmailPartitioningContext(file=f)
|
||||
assert ctx.metadata_last_modified is None
|
||||
|
||||
# -- .msg ------------------------------------
|
||||
|
||||
def it_loads_the_email_message_from_the_filesystem_when_a_path_is_provided(self):
|
||||
ctx = EmailPartitioningContext(file_path=example_doc_path("eml/simple-rfc-822.eml"))
|
||||
assert isinstance(ctx.msg, EmailMessage)
|
||||
|
||||
def and_it_loads_the_email_message_from_a_file_like_object_when_one_is_provided(self):
|
||||
with open(example_doc_path("eml/simple-rfc-822.eml"), "rb") as f:
|
||||
ctx = EmailPartitioningContext(file=f)
|
||||
assert isinstance(ctx.msg, EmailMessage)
|
||||
|
||||
# -- .partitioning_kwargs --------------------
|
||||
|
||||
def it_passes_along_the_kwargs_it_received_on_construction(self):
|
||||
kwargs = {"foo": "bar", "baz": "qux"}
|
||||
ctx = EmailPartitioningContext(kwargs=kwargs)
|
||||
|
||||
assert ctx.partitioning_kwargs == kwargs
|
||||
|
||||
# -- .process_attachments --------------------
|
||||
|
||||
@pytest.mark.parametrize("process_attachments", [True, False])
|
||||
def it_knows_whether_the_caller_wants_to_also_partition_attachments(
|
||||
self, process_attachments: bool
|
||||
):
|
||||
ctx = EmailPartitioningContext(process_attachments=process_attachments)
|
||||
assert ctx.process_attachments == process_attachments
|
||||
|
||||
def but_by_default_it_ignores_attachments(self):
|
||||
ctx = EmailPartitioningContext()
|
||||
assert ctx.process_attachments is False
|
||||
|
||||
# -- .subject --------------------------------
|
||||
|
||||
def it_provides_access_to_the_email_Subject_as_a_string(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/mime-word-encoded-subject.eml"))
|
||||
assert ctx.subject == "Simple email with ☸☿ Unicode subject"
|
||||
|
||||
def but_it_returns_None_when_there_is_no_Subject_header(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/mime-no-subject.eml"))
|
||||
assert ctx.subject is None
|
||||
|
||||
# -- .to_addresses ---------------------------
|
||||
|
||||
def it_provides_access_to_the_To_addresses_when_present(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/mime-multi-to-cc-bcc.eml"))
|
||||
assert ctx.to_addresses == ["Bob <bob@example.com>", "Sue <sue@example.com>"]
|
||||
|
||||
def but_it_returns_None_when_there_are_no_To_addresses(self):
|
||||
ctx = EmailPartitioningContext(example_doc_path("eml/mime-no-to.eml"))
|
||||
assert ctx.to_addresses is None
|
||||
|
||||
# -- fixtures --------------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture()
|
||||
def ctx_args(self) -> dict[str, Any]:
|
||||
return {
|
||||
"file_path": None,
|
||||
"file": None,
|
||||
"content_source": "text/html",
|
||||
"metadata_file_path": None,
|
||||
"metadata_last_modified": None,
|
||||
"process_attachments": False,
|
||||
"kwargs": {},
|
||||
}
|
||||
|
||||
@pytest.fixture()
|
||||
def get_last_modified_date_(self, request: FixtureRequest) -> Mock:
|
||||
return function_mock(request, "unstructured.partition.email.get_last_modified_date")
|
||||
@@ -0,0 +1,174 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.unit_utils import assert_round_trips_through_JSON, example_doc_path
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import Table, Text
|
||||
from unstructured.partition.epub import partition_epub
|
||||
from unstructured.partition.utils.constants import UNSTRUCTURED_INCLUDE_DEBUG_METADATA
|
||||
|
||||
|
||||
def test_partition_epub_from_filename():
|
||||
elements = partition_epub(example_doc_path("simple.epub"))
|
||||
|
||||
assert len(elements) > 0
|
||||
assert isinstance(elements[0], Text)
|
||||
assert elements[1].text.startswith("a shared culture")
|
||||
if UNSTRUCTURED_INCLUDE_DEBUG_METADATA:
|
||||
assert {element.metadata.detection_origin for element in elements} == {"epub"}
|
||||
|
||||
|
||||
def test_partition_epub_from_filename_returns_table_in_elements():
|
||||
elements = partition_epub(example_doc_path("winter-sports.epub"))
|
||||
assert elements[12] == Table(
|
||||
"Contents. List of Illustrations (In certain versions of this etext [in certain\nbrowsers]"
|
||||
" clicking on the image will bring up a larger\nversion.) (etext transcriber's note)"
|
||||
)
|
||||
|
||||
|
||||
def test_partition_epub_from_file():
|
||||
with open(example_doc_path("winter-sports.epub"), "rb") as f:
|
||||
elements = partition_epub(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert elements[2].text.startswith("The Project Gutenberg eBook of Winter Sports")
|
||||
|
||||
|
||||
# -- .metadata.filename --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_epub_from_filename_gets_filename_from_filename_arg():
|
||||
elements = partition_epub(example_doc_path("simple.epub"))
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename == "simple.epub" for e in elements)
|
||||
|
||||
|
||||
def test_partition_epub_from_file_gets_filename_None():
|
||||
with open(example_doc_path("simple.epub"), "rb") as f:
|
||||
elements = partition_epub(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_epub_from_filename_prefers_metadata_filename():
|
||||
elements = partition_epub(example_doc_path("simple.epub"), metadata_filename="orig-name.epub")
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(element.metadata.filename == "orig-name.epub" for element in elements)
|
||||
|
||||
|
||||
def test_partition_epub_from_file_prefers_metadata_filename():
|
||||
with open(example_doc_path("simple.epub"), "rb") as f:
|
||||
elements = partition_epub(file=f, metadata_filename="orig-name.epub")
|
||||
|
||||
assert all(e.metadata.filename == "orig-name.epub" for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.filetype --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_epub_gets_the_EPUB_MIME_type_in_metadata_filetype():
|
||||
EPUB_MIME_TYPE = "application/epub"
|
||||
elements = partition_epub(example_doc_path("simple.epub"))
|
||||
assert all(e.metadata.filetype == EPUB_MIME_TYPE for e in elements), (
|
||||
f"Expected all elements to have '{EPUB_MIME_TYPE}' as their filetype, but got:"
|
||||
f" {repr(elements[0].metadata.filetype)}"
|
||||
)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_epub_from_file_path_gets_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2024-06-14T16:01:29"
|
||||
mocker.patch(
|
||||
"unstructured.partition.epub.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_epub(example_doc_path("winter-sports.epub"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_epub_from_file_gets_last_modified_None():
|
||||
with open(example_doc_path("simple.epub"), "rb") as f:
|
||||
elements = partition_epub(file=f)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_epub_from_file_path_prefers_metadata_last_modified(mocker: MockFixture):
|
||||
filesystem_last_modified = "2024-06-14T16:01:29"
|
||||
metadata_last_modified = "2020-03-08T06:10:23"
|
||||
mocker.patch(
|
||||
"unstructured.partition.epub.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_epub(
|
||||
example_doc_path("winter-sports.epub"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_epub_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-03-08T06:10:23"
|
||||
with open(example_doc_path("simple.epub"), "rb") as f:
|
||||
elements = partition_epub(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified is metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_epub_with_json():
|
||||
filename = "example-docs/winter-sports.epub"
|
||||
elements = partition_epub(filename=filename)
|
||||
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
def test_add_chunking_strategy_on_partition_epub():
|
||||
file_path = example_doc_path("winter-sports.epub")
|
||||
elements = partition_epub(file_path)
|
||||
chunk_elements = partition_epub(file_path, chunking_strategy="by_title")
|
||||
chunks = chunk_by_title(elements)
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_add_chunking_strategy_on_partition_epub_non_default():
|
||||
file_path = example_doc_path("winter-sports.epub")
|
||||
elements = partition_epub(filename=file_path)
|
||||
chunk_elements = partition_epub(
|
||||
file_path,
|
||||
chunking_strategy="by_title",
|
||||
max_characters=5,
|
||||
new_after_n_chars=5,
|
||||
combine_text_under_n_chars=0,
|
||||
)
|
||||
chunks = chunk_by_title(
|
||||
elements,
|
||||
max_characters=5,
|
||||
new_after_n_chars=5,
|
||||
combine_text_under_n_chars=0,
|
||||
)
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_partition_epub_element_metadata_has_languages():
|
||||
filename = example_doc_path("winter-sports.epub")
|
||||
elements = partition_epub(filename=filename)
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_epub_respects_detect_language_per_element():
|
||||
filename = "example-docs/language-docs/eng_spa_mult.epub"
|
||||
elements = partition_epub(filename=filename, detect_language_per_element=True)
|
||||
langs = [element.metadata.languages for element in elements]
|
||||
assert langs == [["eng"], ["spa", "eng"], ["eng"], ["eng"], ["spa"]]
|
||||
@@ -0,0 +1,313 @@
|
||||
"""Test-suite for `unstructured.partition.json` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.unit_utils import example_doc_path
|
||||
from unstructured.documents.elements import CompositeElement
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.email import partition_email
|
||||
from unstructured.partition.html import partition_html
|
||||
from unstructured.partition.json import partition_json
|
||||
from unstructured.partition.text import partition_text
|
||||
from unstructured.partition.xml import partition_xml
|
||||
from unstructured.staging.base import elements_to_json
|
||||
|
||||
DIRECTORY = pathlib.Path(__file__).parent.resolve()
|
||||
|
||||
is_in_docker = os.path.exists("/.dockerenv")
|
||||
|
||||
test_files = [
|
||||
"fake-text.txt",
|
||||
"fake-html.html",
|
||||
"eml/fake-email.eml",
|
||||
]
|
||||
|
||||
is_in_docker = os.path.exists("/.dockerenv")
|
||||
|
||||
|
||||
def test_it_chunks_elements_when_a_chunking_strategy_is_specified():
|
||||
chunks = partition_json(
|
||||
"example-docs/spring-weather.html.json", chunking_strategy="basic", max_characters=1500
|
||||
)
|
||||
|
||||
assert len(chunks) == 9
|
||||
assert all(isinstance(ch, CompositeElement) for ch in chunks)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", test_files)
|
||||
def test_partition_json_from_filename(filename: str):
|
||||
path = example_doc_path(filename)
|
||||
elements = []
|
||||
filetype = FileType.from_extension(os.path.splitext(path)[1])
|
||||
if filetype == FileType.TXT:
|
||||
elements = partition_text(filename=path)
|
||||
if filetype == FileType.HTML:
|
||||
elements = partition_html(filename=path)
|
||||
if filetype == FileType.XML:
|
||||
elements = partition_xml(filename=path)
|
||||
if filetype == FileType.EML:
|
||||
elements = partition_email(filename=path)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
_filename = os.path.basename(filename)
|
||||
test_path = os.path.join(tmpdir, _filename + ".json")
|
||||
elements_to_json(elements, filename=test_path, indent=2)
|
||||
test_elements = partition_json(filename=test_path)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert len(str(elements[0])) > 0
|
||||
|
||||
assert len(elements) == len(test_elements)
|
||||
for i in range(len(elements)):
|
||||
assert elements[i] == test_elements[i]
|
||||
assert elements[i].metadata.filename == filename.split("/")[-1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", test_files)
|
||||
def test_partition_json_from_filename_with_metadata_filename(filename: str):
|
||||
path = example_doc_path(filename)
|
||||
elements = []
|
||||
filetype = FileType.from_extension(os.path.splitext(path)[1])
|
||||
if filetype == FileType.TXT:
|
||||
elements = partition_text(filename=path)
|
||||
if filetype == FileType.HTML:
|
||||
elements = partition_html(filename=path)
|
||||
if filetype == FileType.XML:
|
||||
elements = partition_xml(filename=path)
|
||||
if filetype == FileType.EML:
|
||||
elements = partition_email(filename=path)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
_filename = os.path.basename(filename)
|
||||
test_path = os.path.join(tmpdir, _filename + ".json")
|
||||
elements_to_json(elements, filename=test_path, indent=2)
|
||||
test_elements = partition_json(filename=test_path, metadata_filename="test")
|
||||
|
||||
assert len(test_elements) > 0
|
||||
assert len(str(test_elements[0])) > 0
|
||||
assert all(element.metadata.filename == "test" for element in test_elements)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", test_files)
|
||||
def test_partition_json_from_file(filename: str):
|
||||
path = example_doc_path(filename)
|
||||
elements = []
|
||||
filetype = FileType.from_extension(os.path.splitext(path)[1])
|
||||
if filetype == FileType.TXT:
|
||||
elements = partition_text(filename=path)
|
||||
if filetype == FileType.HTML:
|
||||
elements = partition_html(filename=path)
|
||||
if filetype == FileType.XML:
|
||||
elements = partition_xml(filename=path)
|
||||
if filetype == FileType.EML:
|
||||
elements = partition_email(filename=path)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
_filename = os.path.basename(filename)
|
||||
test_path = os.path.join(tmpdir, _filename + ".json")
|
||||
elements_to_json(elements, filename=test_path, indent=2)
|
||||
with open(test_path, "rb") as f:
|
||||
test_elements = partition_json(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert len(str(elements[0])) > 0
|
||||
assert len(elements) == len(test_elements)
|
||||
for i in range(len(elements)):
|
||||
assert elements[i] == test_elements[i]
|
||||
assert elements[i].metadata.filename == filename.split("/")[-1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", test_files)
|
||||
def test_partition_json_from_file_with_metadata_filename(filename: str):
|
||||
path = example_doc_path(filename)
|
||||
elements = []
|
||||
filetype = FileType.from_extension(os.path.splitext(path)[1])
|
||||
if filetype == FileType.TXT:
|
||||
elements = partition_text(filename=path)
|
||||
if filetype == FileType.HTML:
|
||||
elements = partition_html(filename=path)
|
||||
if filetype == FileType.XML:
|
||||
elements = partition_xml(filename=path)
|
||||
if filetype == FileType.EML:
|
||||
elements = partition_email(filename=path)
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
_filename = os.path.basename(filename)
|
||||
test_path = os.path.join(tmpdir, _filename + ".json")
|
||||
elements_to_json(elements, filename=test_path, indent=2)
|
||||
with open(test_path, "rb") as f:
|
||||
test_elements = partition_json(file=f, metadata_filename="test")
|
||||
|
||||
for i in range(len(test_elements)):
|
||||
assert test_elements[i].metadata.filename == "test"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", test_files)
|
||||
def test_partition_json_from_text(filename: str):
|
||||
path = example_doc_path(filename)
|
||||
elements = []
|
||||
filetype = FileType.from_extension(os.path.splitext(path)[1])
|
||||
if filetype == FileType.TXT:
|
||||
elements = partition_text(filename=path)
|
||||
if filetype == FileType.HTML:
|
||||
elements = partition_html(filename=path)
|
||||
if filetype == FileType.XML:
|
||||
elements = partition_xml(filename=path)
|
||||
if filetype == FileType.EML:
|
||||
elements = partition_email(filename=path)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
_filename = os.path.basename(filename)
|
||||
test_path = os.path.join(tmpdir, _filename + ".json")
|
||||
elements_to_json(elements, filename=test_path, indent=2)
|
||||
with open(test_path) as f:
|
||||
text = f.read()
|
||||
test_elements = partition_json(text=text)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert len(str(elements[0])) > 0
|
||||
assert len(elements) == len(test_elements)
|
||||
for i in range(len(elements)):
|
||||
assert elements[i] == test_elements[i]
|
||||
assert elements[i].metadata.filename == filename.split("/")[-1]
|
||||
|
||||
|
||||
def test_partition_json_raises_with_none_specified():
|
||||
with pytest.raises(ValueError):
|
||||
partition_json()
|
||||
|
||||
|
||||
def test_partition_json_works_with_empty_string():
|
||||
assert partition_json(text="") == []
|
||||
|
||||
|
||||
def test_partition_json_fails_with_empty_item():
|
||||
with pytest.raises(ValueError):
|
||||
partition_json(text="{}")
|
||||
|
||||
|
||||
def test_partition_json_works_with_empty_list():
|
||||
assert partition_json(text="[]") == []
|
||||
|
||||
|
||||
def test_partition_json_raises_with_too_many_specified():
|
||||
path = example_doc_path("fake-text.txt")
|
||||
elements = []
|
||||
filetype = FileType.from_extension(os.path.splitext(path)[1])
|
||||
if filetype == FileType.TXT:
|
||||
elements = partition_text(filename=path)
|
||||
if filetype == FileType.HTML:
|
||||
elements = partition_html(filename=path)
|
||||
if filetype == FileType.XML:
|
||||
elements = partition_xml(filename=path)
|
||||
if filetype == FileType.EML:
|
||||
elements = partition_email(filename=path)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
test_path = os.path.join(tmpdir, "fake-text.txt.json")
|
||||
elements_to_json(elements, filename=test_path, indent=2)
|
||||
with open(test_path, "rb") as f:
|
||||
text = f.read().decode("utf-8")
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_json(filename=test_path, file=f)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_json(filename=test_path, text=text)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_json(file=f, text=text)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_json(filename=test_path, file=f, text=text)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_json_from_file_path_gets_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.json.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_json(example_doc_path("spring-weather.html.json"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_json_from_file_gets_last_modified_None():
|
||||
with open("example-docs/spring-weather.html.json", "rb") as f:
|
||||
elements = partition_json(file=f)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_json_from_text_gets_last_modified_None():
|
||||
with open("example-docs/spring-weather.html.json") as f:
|
||||
text = f.read()
|
||||
|
||||
elements = partition_json(text=text)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_json_from_file_path_prefers_metadata_last_modified(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.json.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_json(
|
||||
"example-docs/spring-weather.html.json", metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_json_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
with open(example_doc_path("spring-weather.html.json"), "rb") as f:
|
||||
elements = partition_json(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_json_from_text_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
with open("example-docs/spring-weather.html.json") as f:
|
||||
text = f.read()
|
||||
|
||||
elements = partition_json(text=text, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_json_raises_with_unprocessable_json_array():
|
||||
text = '[{"invalid": "schema"}]'
|
||||
with pytest.raises(ValueError):
|
||||
partition_json(text=text)
|
||||
|
||||
|
||||
def test_partition_json_raises_with_unprocessable_json():
|
||||
# NOTE(robinson) - This is unprocessable because it is not a list of dicts,
|
||||
# per the Unstructured ISD format
|
||||
text = '{"hi": "there"}'
|
||||
with pytest.raises(ValueError):
|
||||
partition_json(text=text)
|
||||
|
||||
|
||||
def test_partition_json_raises_with_invalid_json():
|
||||
text = '[{"hi": "there"}]]'
|
||||
with pytest.raises(ValueError):
|
||||
partition_json(text=text)
|
||||
@@ -0,0 +1,334 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.unit_utils import assert_round_trips_through_JSON, example_doc_path
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import ElementType, Title
|
||||
from unstructured.partition.md import partition_md
|
||||
from unstructured.partition.utils.constants import UNSTRUCTURED_INCLUDE_DEBUG_METADATA
|
||||
|
||||
|
||||
def test_partition_md_from_filename():
|
||||
filename = example_doc_path("README.md")
|
||||
elements = partition_md(filename=filename)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert "PageBreak" not in [elem.category for elem in elements]
|
||||
assert isinstance(elements[0], Title)
|
||||
if UNSTRUCTURED_INCLUDE_DEBUG_METADATA:
|
||||
assert {element.metadata.detection_origin for element in elements} == {"md"}
|
||||
|
||||
|
||||
def test_partition_md_from_file():
|
||||
filename = example_doc_path("README.md")
|
||||
with open(filename, "rb") as f:
|
||||
elements = partition_md(file=f)
|
||||
assert len(elements) > 0
|
||||
|
||||
|
||||
def test_partition_md_from_text():
|
||||
with open(example_doc_path("README.md")) as f:
|
||||
text = f.read()
|
||||
|
||||
elements = partition_md(text=text)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
|
||||
|
||||
class MockResponse:
|
||||
def __init__(self, text: str, status_code: int, headers: dict[str, Any] = {}):
|
||||
self.text = text
|
||||
self.status_code = status_code
|
||||
self.ok = status_code < 300
|
||||
self.headers = headers
|
||||
|
||||
|
||||
def test_partition_md_from_url():
|
||||
filename = example_doc_path("README.md")
|
||||
with open(filename) as f:
|
||||
text = f.read()
|
||||
|
||||
response = MockResponse(
|
||||
text=text,
|
||||
status_code=200,
|
||||
headers={"Content-Type": "text/markdown"},
|
||||
)
|
||||
with patch.object(requests, "get", return_value=response) as _:
|
||||
elements = partition_md(url="https://fake.url")
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_md_from_url_raises_with_bad_status_code():
|
||||
filename = example_doc_path("README.md")
|
||||
with open(filename) as f:
|
||||
text = f.read()
|
||||
|
||||
response = MockResponse(
|
||||
text=text,
|
||||
status_code=500,
|
||||
headers={"Content-Type": "text/html"},
|
||||
)
|
||||
with patch.object(requests, "get", return_value=response) as _, pytest.raises(ValueError):
|
||||
partition_md(url="https://fake.url")
|
||||
|
||||
|
||||
def test_partition_md_from_url_raises_with_bad_content_type():
|
||||
filename = example_doc_path("README.md")
|
||||
with open(filename) as f:
|
||||
text = f.read()
|
||||
|
||||
response = MockResponse(
|
||||
text=text,
|
||||
status_code=200,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
with patch.object(requests, "get", return_value=response) as _, pytest.raises(ValueError):
|
||||
partition_md(url="https://fake.url")
|
||||
|
||||
|
||||
def test_partition_md_raises_with_none_specified():
|
||||
with pytest.raises(ValueError):
|
||||
partition_md()
|
||||
|
||||
|
||||
def test_partition_md_raises_with_too_many_specified():
|
||||
filename = example_doc_path("README.md")
|
||||
with open(filename) as f:
|
||||
text = f.read()
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_md(filename=filename, text=text)
|
||||
|
||||
|
||||
# -- .metadata.filename --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_md_from_filename_gets_filename_from_filename_arg():
|
||||
elements = partition_md(example_doc_path("README.md"))
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename == "README.md" for e in elements)
|
||||
|
||||
|
||||
def test_partition_md_from_file_gets_filename_None():
|
||||
with open(example_doc_path("README.md"), "rb") as f:
|
||||
elements = partition_md(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_md_from_filename_prefers_metadata_filename():
|
||||
elements = partition_md(example_doc_path("README.md"), metadata_filename="orig-name.md")
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(element.metadata.filename == "orig-name.md" for element in elements)
|
||||
|
||||
|
||||
def test_partition_md_from_file_prefers_metadata_filename():
|
||||
with open(example_doc_path("README.md"), "rb") as f:
|
||||
elements = partition_md(file=f, metadata_filename="orig-name.md")
|
||||
|
||||
assert all(e.metadata.filename == "orig-name.md" for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.filetype --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_md_gets_the_MD_MIME_type_in_metadata_filetype():
|
||||
MD_MIME_TYPE = "text/markdown"
|
||||
elements = partition_md(example_doc_path("README.md"))
|
||||
assert all(e.metadata.filetype == MD_MIME_TYPE for e in elements), (
|
||||
f"Expected all elements to have '{MD_MIME_TYPE}' as their filetype, but got:"
|
||||
f" {repr(elements[0].metadata.filetype)}"
|
||||
)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_md_from_file_path_gets_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.md.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_md(example_doc_path("README.md"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_md_from_file_gets_last_modified_None():
|
||||
with open(example_doc_path("README.md"), "rb") as f:
|
||||
elements = partition_md(file=f)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_md_from_text_gets_last_modified_None():
|
||||
with open(example_doc_path("README.md")) as f:
|
||||
text = f.read()
|
||||
|
||||
elements = partition_md(text=text)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_md_from_file_path_prefers_metadata_last_modified(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.md.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_md(
|
||||
example_doc_path("README.md"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_md_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
with open(example_doc_path("README.md"), "rb") as f:
|
||||
elements = partition_md(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_md_from_text_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
with open(example_doc_path("README.md")) as f:
|
||||
text = f.read()
|
||||
|
||||
elements = partition_md(text=text, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_md_with_json():
|
||||
with open(example_doc_path("README.md")) as f:
|
||||
text = f.read()
|
||||
elements = partition_md(text=text)
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
def test_add_chunking_strategy_by_title_on_partition_md():
|
||||
filename = example_doc_path("README.md")
|
||||
elements = partition_md(filename)
|
||||
chunk_elements = partition_md(filename, chunking_strategy="by_title")
|
||||
chunks = chunk_by_title(elements)
|
||||
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_partition_md_element_metadata_has_languages():
|
||||
filename = "example-docs/README.md"
|
||||
elements = partition_md(filename=filename)
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_md_respects_detect_language_per_element():
|
||||
filename = "example-docs/language-docs/eng_spa_mult.md"
|
||||
elements = partition_md(filename=filename, detect_language_per_element=True)
|
||||
langs = [element.metadata.languages for element in elements]
|
||||
assert langs == [["eng"], ["spa", "eng"], ["eng"], ["eng"], ["spa"]]
|
||||
|
||||
|
||||
def test_partition_md_parse_table():
|
||||
filename = example_doc_path("simple-table.md")
|
||||
elements = partition_md(filename=filename)
|
||||
assert len(elements) > 0
|
||||
assert elements[0].category == ElementType.TABLE
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", ["umlauts-utf8.md", "umlauts-non-utf8.md"])
|
||||
def test_partition_md_with_umlauts(filename: str):
|
||||
filename = example_doc_path(filename)
|
||||
elements = partition_md(filename=filename)
|
||||
assert len(elements) > 0
|
||||
assert elements[-1].text.endswith("äöüß")
|
||||
|
||||
|
||||
def test_partition_md_xml_processing_instruction():
|
||||
xml_content = """```
|
||||
<?xml version="1.0"?>
|
||||
<sparql xmlns="http://www.w3.org/2005/sparql-results#">
|
||||
<head></head>
|
||||
<boolean>true</boolean>
|
||||
</sparql>
|
||||
```"""
|
||||
|
||||
elements = partition_md(text=xml_content)
|
||||
assert len(elements) == 1
|
||||
|
||||
|
||||
def test_partition_md_xml_processing_instruction_with_indents():
|
||||
xml_content = """```
|
||||
<?xml version="1.0"?>
|
||||
<sparql xmlns="http://www.w3.org/2005/sparql-results#">
|
||||
<head></head>
|
||||
<boolean>true</boolean>
|
||||
</sparql>
|
||||
```"""
|
||||
|
||||
elements = partition_md(text=xml_content)
|
||||
assert len(elements) == 1
|
||||
|
||||
|
||||
def test_partition_md_non_xml_processing_instruction():
|
||||
php_content = """```
|
||||
<?php echo "hello"; ?>
|
||||
```"""
|
||||
|
||||
elements = partition_md(text=php_content)
|
||||
assert len(elements) == 1
|
||||
|
||||
|
||||
def test_partition_fenced_code():
|
||||
filename = example_doc_path("codeblock.md")
|
||||
elements = partition_md(filename=filename)
|
||||
|
||||
# Should have 5 elements: 2 titles and 3 code blocks
|
||||
assert len(elements) == 5
|
||||
|
||||
assert elements[0].text == "HTML Example"
|
||||
|
||||
expected_html = """<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Sample HTML</title>
|
||||
</head>
|
||||
<body>
|
||||
<h1>Hello, World!</h1>
|
||||
<p>This is a simple HTML example.</p>
|
||||
</body>
|
||||
</html>"""
|
||||
assert elements[1].text == expected_html
|
||||
|
||||
assert elements[2].text == "XML Example"
|
||||
|
||||
expected_xml = """<note>
|
||||
<to>Tove</to>
|
||||
<from>Jani</from>
|
||||
<heading>Reminder</heading>
|
||||
<body>Don't forget me this weekend!</body>
|
||||
</note>"""
|
||||
assert elements[3].text == expected_xml
|
||||
|
||||
assert elements[4].text == expected_xml
|
||||
@@ -0,0 +1,614 @@
|
||||
"""Test suite for `unstructured.partition.msg` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from oxmsg import Message
|
||||
|
||||
from test_unstructured.unit_utils import (
|
||||
FixtureRequest,
|
||||
LogCaptureFixture,
|
||||
Mock,
|
||||
assert_round_trips_through_JSON,
|
||||
example_doc_path,
|
||||
function_mock,
|
||||
property_mock,
|
||||
)
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import (
|
||||
ElementMetadata,
|
||||
ListItem,
|
||||
NarrativeText,
|
||||
Text,
|
||||
)
|
||||
from unstructured.partition.common import UnsupportedFileFormatError
|
||||
from unstructured.partition.msg import MsgPartitionerOptions, partition_msg
|
||||
|
||||
EXPECTED_MSG_OUTPUT = [
|
||||
NarrativeText(text="This is a test email to use for unit tests."),
|
||||
Text(text="Important points:"),
|
||||
ListItem(text="Roses are red"),
|
||||
ListItem(text="Violets are blue"),
|
||||
]
|
||||
|
||||
|
||||
def test_partition_msg_from_filename():
|
||||
filename = example_doc_path("fake-email.msg")
|
||||
elements = partition_msg(filename=filename)
|
||||
parent_id = elements[0].metadata.parent_id
|
||||
|
||||
assert elements == EXPECTED_MSG_OUTPUT
|
||||
assert (
|
||||
elements[0].metadata.to_dict()
|
||||
== ElementMetadata(
|
||||
coordinates=None,
|
||||
filename=filename,
|
||||
last_modified="2023-03-28T17:00:31+00:00",
|
||||
page_number=None,
|
||||
url=None,
|
||||
sent_from=['"Matthew Robinson" <mrobinson@unstructured.io>'],
|
||||
sent_to=["mrobinson@unstructured.io"],
|
||||
subject="Test Email",
|
||||
filetype="application/vnd.ms-outlook",
|
||||
parent_id=parent_id,
|
||||
languages=["eng"],
|
||||
).to_dict()
|
||||
)
|
||||
|
||||
|
||||
def test_partition_msg_from_filename_returns_uns_elements():
|
||||
filename = example_doc_path("fake-email.msg")
|
||||
elements = partition_msg(filename=filename)
|
||||
assert isinstance(elements[0], NarrativeText)
|
||||
|
||||
|
||||
def test_partition_msg_from_filename_with_metadata_filename():
|
||||
filename = example_doc_path("fake-email.msg")
|
||||
elements = partition_msg(filename=filename, metadata_filename="test")
|
||||
assert all(element.metadata.filename == "test" for element in elements)
|
||||
|
||||
|
||||
def test_partition_msg_from_filename_with_text_content():
|
||||
filename = example_doc_path("fake-email.msg")
|
||||
|
||||
elements = partition_msg(filename=filename)
|
||||
|
||||
assert str(elements[0]) == "This is a test email to use for unit tests."
|
||||
assert elements[0].metadata.filename == "fake-email.msg"
|
||||
assert elements[0].metadata.file_directory == example_doc_path("")
|
||||
|
||||
|
||||
def test_partition_msg_raises_with_missing_file():
|
||||
filename = example_doc_path("doesnt-exist.msg")
|
||||
with pytest.raises(FileNotFoundError):
|
||||
partition_msg(filename=filename)
|
||||
|
||||
|
||||
def test_partition_msg_from_file():
|
||||
filename = example_doc_path("fake-email.msg")
|
||||
with open(filename, "rb") as f:
|
||||
elements = partition_msg(file=f)
|
||||
assert elements == EXPECTED_MSG_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename is None
|
||||
|
||||
|
||||
def test_partition_msg_from_file_with_metadata_filename():
|
||||
filename = example_doc_path("fake-email.msg")
|
||||
with open(filename, "rb") as f:
|
||||
elements = partition_msg(file=f, metadata_filename="test")
|
||||
assert elements == EXPECTED_MSG_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "test"
|
||||
|
||||
|
||||
def test_partition_msg_uses_file_path_when_both_are_specified():
|
||||
elements = partition_msg(example_doc_path("fake-email.msg"), file=io.BytesIO(b"abcde"))
|
||||
assert elements == EXPECTED_MSG_OUTPUT
|
||||
|
||||
|
||||
def test_partition_msg_raises_with_neither():
|
||||
with pytest.raises(ValueError):
|
||||
partition_msg()
|
||||
|
||||
|
||||
# -- attachments ---------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_msg_can_process_attachments():
|
||||
elements = partition_msg(
|
||||
example_doc_path("fake-email-multiple-attachments.msg"), process_attachments=True
|
||||
)
|
||||
|
||||
assert all(e.metadata.filename == "fake-email-multiple-attachments.msg" for e in elements[:5])
|
||||
assert all(e.metadata.filename == "unstructured_logo.png" for e in elements[5:7])
|
||||
assert all(e.metadata.filename == "dense_doc.pdf" for e in elements[7:343])
|
||||
assert all(e.metadata.filename == "Engineering Onboarding.pptx" for e in elements[343:])
|
||||
assert [e.text for e in elements[:5]] == [
|
||||
"Here are those documents.",
|
||||
"--",
|
||||
"Mallori Harrell",
|
||||
"Unstructured Technologies",
|
||||
"Data Scientist",
|
||||
]
|
||||
assert [type(e).__name__ for e in elements][:10] == [
|
||||
"NarrativeText",
|
||||
"Text",
|
||||
"Text",
|
||||
"Text",
|
||||
"Text",
|
||||
"Image",
|
||||
"Text",
|
||||
"Text",
|
||||
"Title",
|
||||
"Title",
|
||||
]
|
||||
assert [type(e).__name__ for e in elements][-10:] == [
|
||||
"Title",
|
||||
"ListItem",
|
||||
"ListItem",
|
||||
"ListItem",
|
||||
"ListItem",
|
||||
"ListItem",
|
||||
"ListItem",
|
||||
"ListItem",
|
||||
"ListItem",
|
||||
"ListItem",
|
||||
]
|
||||
|
||||
|
||||
def test_partition_msg_silently_skips_attachments_it_cannot_partition(request: FixtureRequest):
|
||||
function_mock(
|
||||
request, "unstructured.partition.auto.partition", side_effect=UnsupportedFileFormatError()
|
||||
)
|
||||
|
||||
elements = partition_msg(
|
||||
example_doc_path("fake-email-multiple-attachments.msg"), process_attachments=True
|
||||
)
|
||||
|
||||
# -- no exception is raised --
|
||||
assert elements == [
|
||||
# -- the email body is partitioned --
|
||||
NarrativeText("Here are those documents."),
|
||||
Text("--"),
|
||||
Text("Mallori Harrell"),
|
||||
Text("Unstructured Technologies"),
|
||||
Text("Data Scientist"),
|
||||
# -- no elements appear for the attachment(s) --
|
||||
]
|
||||
|
||||
|
||||
# -- .metadata.filename --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_msg_from_filename_gets_filename_metadata_from_file_path():
|
||||
elements = partition_msg(example_doc_path("fake-email.msg"))
|
||||
|
||||
assert all(e.metadata.filename == "fake-email.msg" for e in elements)
|
||||
assert all(e.metadata.file_directory == example_doc_path("") for e in elements)
|
||||
|
||||
|
||||
def test_partition_msg_from_file_gets_filename_metadata_None():
|
||||
with open(example_doc_path("fake-email.msg"), "rb") as f:
|
||||
elements = partition_msg(file=f)
|
||||
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
assert all(e.metadata.file_directory is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_msg_from_filename_prefers_metadata_filename():
|
||||
elements = partition_msg(example_doc_path("fake-email.msg"), metadata_filename="a/b/c.msg")
|
||||
|
||||
assert all(e.metadata.filename == "c.msg" for e in elements)
|
||||
assert all(e.metadata.file_directory == "a/b" for e in elements)
|
||||
|
||||
|
||||
def test_partition_msg_from_file_prefers_metadata_filename():
|
||||
with open(example_doc_path("fake-email.msg"), "rb") as f:
|
||||
elements = partition_msg(file=f, metadata_filename="d/e/f.msg")
|
||||
|
||||
assert all(e.metadata.filename == "f.msg" for e in elements)
|
||||
assert all(e.metadata.file_directory == "d/e" for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.filetype --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_msg_gets_the_MSG_mime_type_in_metadata_filetype():
|
||||
MSG_MIME_TYPE = "application/vnd.ms-outlook"
|
||||
elements = partition_msg(example_doc_path("fake-email.msg"))
|
||||
assert all(e.metadata.filetype == MSG_MIME_TYPE for e in elements), (
|
||||
f"Expected all elements to have '{MSG_MIME_TYPE}' as their filetype, but got:"
|
||||
f" {repr(elements[0].metadata.filetype)}"
|
||||
)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_msg_pulls_last_modified_from_message_sent_date():
|
||||
elements = partition_msg(example_doc_path("fake-email.msg"))
|
||||
assert all(e.metadata.last_modified == "2023-03-28T17:00:31+00:00" for e in elements)
|
||||
|
||||
|
||||
def test_partition_msg_from_file_path_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
|
||||
elements = partition_msg(
|
||||
example_doc_path("fake-email.msg"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert elements[0].metadata.last_modified == metadata_last_modified
|
||||
|
||||
|
||||
def test_partition_msg_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
|
||||
with open(example_doc_path("fake-email.msg"), "rb") as f:
|
||||
elements = partition_msg(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_msg_with_json():
|
||||
elements = partition_msg(example_doc_path("fake-email.msg"))
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
def test_partition_msg_with_pgp_encrypted_message(caplog: LogCaptureFixture):
|
||||
elements = partition_msg(example_doc_path("fake-encrypted.msg"))
|
||||
|
||||
assert elements == []
|
||||
assert "WARNING" in caplog.text
|
||||
assert "Encrypted email detected" in caplog.text
|
||||
|
||||
|
||||
def test_add_chunking_strategy_by_title_on_partition_msg():
|
||||
filename = example_doc_path("fake-email.msg")
|
||||
|
||||
elements = partition_msg(filename=filename)
|
||||
chunk_elements = partition_msg(filename, chunking_strategy="by_title")
|
||||
chunks = chunk_by_title(elements)
|
||||
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
# -- language behaviors --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_msg_element_metadata_has_languages():
|
||||
filename = "example-docs/fake-email.msg"
|
||||
elements = partition_msg(filename=filename)
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_msg_respects_languages_arg():
|
||||
filename = "example-docs/fake-email.msg"
|
||||
elements = partition_msg(filename=filename, languages=["deu"])
|
||||
assert all(element.metadata.languages == ["deu"] for element in elements)
|
||||
|
||||
|
||||
def test_partition_msg_raises_TypeError_for_invalid_languages():
|
||||
with pytest.raises(TypeError):
|
||||
filename = "example-docs/fake-email.msg"
|
||||
partition_msg(filename=filename, languages="eng")
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# ISOLATED UNIT TESTS
|
||||
# ================================================================================================
|
||||
# These test components used by `partition_msg()` in isolation such that all edge cases can be
|
||||
# exercised.
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class DescribeMsgAttachmentFilenameSanitization:
|
||||
"""Unit-test suite for filename sanitization in MSG attachments (GHSA-gm8q-m8mv-jj5m)."""
|
||||
|
||||
def it_sanitizes_path_traversal_attempts(self, request: FixtureRequest):
|
||||
from unstructured.partition.msg import _AttachmentPartitioner
|
||||
|
||||
attachment = Mock()
|
||||
attachment.file_name = "../../../etc/passwd"
|
||||
attachment.file_bytes = b"malicious content"
|
||||
attachment.last_modified = None
|
||||
|
||||
opts = Mock()
|
||||
opts.metadata_last_modified = None
|
||||
|
||||
partitioner = _AttachmentPartitioner(attachment, opts)
|
||||
|
||||
assert partitioner._attachment_file_name == "passwd"
|
||||
|
||||
def it_sanitizes_absolute_unix_paths(self, request: FixtureRequest):
|
||||
from unstructured.partition.msg import _AttachmentPartitioner
|
||||
|
||||
attachment = Mock()
|
||||
attachment.file_name = "/etc/passwd"
|
||||
attachment.file_bytes = b"malicious content"
|
||||
attachment.last_modified = None
|
||||
|
||||
opts = Mock()
|
||||
opts.metadata_last_modified = None
|
||||
|
||||
partitioner = _AttachmentPartitioner(attachment, opts)
|
||||
|
||||
assert partitioner._attachment_file_name == "passwd"
|
||||
|
||||
def it_sanitizes_absolute_windows_paths(self, request: FixtureRequest):
|
||||
from unstructured.partition.msg import _AttachmentPartitioner
|
||||
|
||||
attachment = Mock()
|
||||
attachment.file_name = "C:\\Windows\\System32\\config\\sam"
|
||||
attachment.file_bytes = b"malicious content"
|
||||
attachment.last_modified = None
|
||||
|
||||
opts = Mock()
|
||||
opts.metadata_last_modified = None
|
||||
|
||||
partitioner = _AttachmentPartitioner(attachment, opts)
|
||||
|
||||
assert partitioner._attachment_file_name == "sam"
|
||||
|
||||
def it_removes_null_bytes_from_filenames(self, request: FixtureRequest):
|
||||
from unstructured.partition.msg import _AttachmentPartitioner
|
||||
|
||||
attachment = Mock()
|
||||
attachment.file_name = "file\x00.txt"
|
||||
attachment.file_bytes = b"content"
|
||||
attachment.last_modified = None
|
||||
|
||||
opts = Mock()
|
||||
opts.metadata_last_modified = None
|
||||
|
||||
partitioner = _AttachmentPartitioner(attachment, opts)
|
||||
|
||||
assert partitioner._attachment_file_name == "file.txt"
|
||||
assert "\x00" not in partitioner._attachment_file_name
|
||||
|
||||
def it_handles_dot_and_dotdot_filenames(self, request: FixtureRequest):
|
||||
from unstructured.partition.msg import _AttachmentPartitioner
|
||||
|
||||
opts = Mock()
|
||||
opts.metadata_last_modified = None
|
||||
|
||||
# Test single dot
|
||||
attachment1 = Mock()
|
||||
attachment1.file_name = "."
|
||||
attachment1.file_bytes = b"content"
|
||||
attachment1.last_modified = None
|
||||
partitioner1 = _AttachmentPartitioner(attachment1, opts)
|
||||
assert partitioner1._attachment_file_name == "unknown"
|
||||
|
||||
# Test double dot
|
||||
attachment2 = Mock()
|
||||
attachment2.file_name = ".."
|
||||
attachment2.file_bytes = b"content"
|
||||
attachment2.last_modified = None
|
||||
partitioner2 = _AttachmentPartitioner(attachment2, opts)
|
||||
assert partitioner2._attachment_file_name == "unknown"
|
||||
|
||||
def it_handles_missing_filename(self, request: FixtureRequest):
|
||||
from unstructured.partition.msg import _AttachmentPartitioner
|
||||
|
||||
attachment = Mock()
|
||||
attachment.file_name = None
|
||||
attachment.file_bytes = b"content"
|
||||
attachment.last_modified = None
|
||||
|
||||
opts = Mock()
|
||||
opts.metadata_last_modified = None
|
||||
|
||||
partitioner = _AttachmentPartitioner(attachment, opts)
|
||||
|
||||
assert partitioner._attachment_file_name == "unknown"
|
||||
|
||||
def it_allows_valid_filenames_through(self, request: FixtureRequest):
|
||||
from unstructured.partition.msg import _AttachmentPartitioner
|
||||
|
||||
attachment = Mock()
|
||||
attachment.file_name = "document.pdf"
|
||||
attachment.file_bytes = b"content"
|
||||
attachment.last_modified = None
|
||||
|
||||
opts = Mock()
|
||||
opts.metadata_last_modified = None
|
||||
|
||||
partitioner = _AttachmentPartitioner(attachment, opts)
|
||||
|
||||
assert partitioner._attachment_file_name == "document.pdf"
|
||||
|
||||
def it_handles_complex_path_traversal_with_mixed_separators(self, request: FixtureRequest):
|
||||
from unstructured.partition.msg import _AttachmentPartitioner
|
||||
|
||||
attachment = Mock()
|
||||
attachment.file_name = "..\\../\\..\\etc/passwd"
|
||||
attachment.file_bytes = b"malicious content"
|
||||
attachment.last_modified = None
|
||||
|
||||
opts = Mock()
|
||||
opts.metadata_last_modified = None
|
||||
|
||||
partitioner = _AttachmentPartitioner(attachment, opts)
|
||||
|
||||
assert partitioner._attachment_file_name == "passwd"
|
||||
|
||||
def it_handles_empty_string_filename(self, request: FixtureRequest):
|
||||
from unstructured.partition.msg import _AttachmentPartitioner
|
||||
|
||||
attachment = Mock()
|
||||
attachment.file_name = ""
|
||||
attachment.file_bytes = b"content"
|
||||
attachment.last_modified = None
|
||||
|
||||
opts = Mock()
|
||||
opts.metadata_last_modified = None
|
||||
|
||||
partitioner = _AttachmentPartitioner(attachment, opts)
|
||||
|
||||
assert partitioner._attachment_file_name == "unknown"
|
||||
|
||||
|
||||
class DescribeMsgPartitionerOptions:
|
||||
"""Unit-test suite for `unstructured.partition.msg.MsgPartitionerOptions` objects."""
|
||||
|
||||
# -- .extra_msg_metadata ---------------------
|
||||
|
||||
def it_provides_email_specific_metadata_to_add_to_each_element(self, opts_args: dict[str, Any]):
|
||||
opts_args["file_path"] = example_doc_path("fake-email-with-cc-and-bcc.msg")
|
||||
opts = MsgPartitionerOptions(**opts_args)
|
||||
|
||||
m = opts.extra_msg_metadata
|
||||
assert m.bcc_recipient == ["hello@unstructured.io"]
|
||||
assert m.cc_recipient == ["steve@unstructured.io"]
|
||||
assert m.email_message_id == "14DDEF33-2BA7-4CDD-A4D8-E7C5873B37F2@gmail.com"
|
||||
assert m.sent_from == ['"John" <johnjennings702@gmail.com>']
|
||||
assert m.sent_to == [
|
||||
"john-ctr@unstructured.io",
|
||||
"steve@unstructured.io",
|
||||
"hello@unstructured.io",
|
||||
]
|
||||
assert m.subject == "Fake email with cc and bcc recipients"
|
||||
|
||||
# -- .is_encrypted ---------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_name", "expected_value"), [("fake-encrypted.msg", True), ("fake-email.msg", False)]
|
||||
)
|
||||
def it_knows_when_the_msg_is_encrypted(
|
||||
self, file_name: str, expected_value: bool, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["file_path"] = example_doc_path(file_name)
|
||||
opts = MsgPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.is_encrypted is expected_value
|
||||
|
||||
# -- .metadata_file_path ---------------------
|
||||
|
||||
def it_uses_the_metadata_file_path_arg_when_provided(self, opts_args: dict[str, Any]):
|
||||
opts_args["file_path"] = "x/y/z.msg"
|
||||
opts_args["metadata_file_path"] = "a/b/c.msg"
|
||||
opts = MsgPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.metadata_file_path == "a/b/c.msg"
|
||||
|
||||
def and_it_falls_back_to_the_MSG_file_path_arg_when_provided(self, opts_args: dict[str, Any]):
|
||||
file_path = example_doc_path("fake-email.msg")
|
||||
opts_args["file_path"] = file_path
|
||||
opts = MsgPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.metadata_file_path == file_path
|
||||
|
||||
def but_it_returns_None_when_neither_path_is_available(self, opts_args: dict[str, Any]):
|
||||
opts = MsgPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.metadata_file_path is None
|
||||
|
||||
# -- .metadata_last_modified -----------------
|
||||
|
||||
def it_uses_metadata_last_modified_when_provided_by_the_caller(self, opts_args: dict[str, Any]):
|
||||
metadata_last_modified = "2024-03-05T17:02:53"
|
||||
opts_args["metadata_last_modified"] = metadata_last_modified
|
||||
opts_args["file_path"] = example_doc_path("fake-email.msg")
|
||||
opts = MsgPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.metadata_last_modified == metadata_last_modified
|
||||
|
||||
def and_it_uses_the_message_Date_header_when_metadata_last_modified_is_not_provided(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["file_path"] = example_doc_path("fake-email.msg")
|
||||
opts = MsgPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.metadata_last_modified == "2023-03-28T17:00:31+00:00"
|
||||
|
||||
@pytest.mark.parametrize("filesystem_last_modified", ["2024-06-03T20:12:53", None])
|
||||
def and_it_uses_the_last_modified_date_from_the_source_file_when_the_message_has_no_sent_date(
|
||||
self,
|
||||
opts_args: dict[str, Any],
|
||||
filesystem_last_modified: str | None,
|
||||
Message_sent_date_: Mock,
|
||||
_last_modified_prop_: Mock,
|
||||
):
|
||||
Message_sent_date_.return_value = None
|
||||
_last_modified_prop_.return_value = filesystem_last_modified
|
||||
opts_args["file_path"] = example_doc_path("fake-email.msg")
|
||||
opts = MsgPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.metadata_last_modified == filesystem_last_modified
|
||||
|
||||
# -- .msg ------------------------------------
|
||||
|
||||
def it_loads_the_msg_document_from_a_file_path_when_provided(self, opts_args: dict[str, Any]):
|
||||
opts_args["file_path"] = example_doc_path("fake-email.msg")
|
||||
opts = MsgPartitionerOptions(**opts_args)
|
||||
|
||||
assert isinstance(opts.msg, Message)
|
||||
|
||||
def and_it_loads_the_msg_document_from_a_file_like_object_when_provided(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
with open(example_doc_path("fake-email.msg"), "rb") as f:
|
||||
opts_args["file"] = io.BytesIO(f.read())
|
||||
opts = MsgPartitionerOptions(**opts_args)
|
||||
|
||||
assert isinstance(opts.msg, Message)
|
||||
|
||||
def but_it_raises_when_neither_is_provided(self, opts_args: dict[str, Any]):
|
||||
with pytest.raises(ValueError, match="one of `file` or `filename` arguments must be prov"):
|
||||
MsgPartitionerOptions(**opts_args).msg
|
||||
|
||||
# -- .partition_attachments ------------------
|
||||
|
||||
@pytest.mark.parametrize("partition_attachments", [True, False])
|
||||
def it_knows_whether_attachments_should_also_be_partitioned(
|
||||
self, partition_attachments: bool, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["file_path"] = example_doc_path("fake-email.msg")
|
||||
opts_args["partition_attachments"] = partition_attachments
|
||||
opts = MsgPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.partition_attachments is partition_attachments
|
||||
|
||||
# -- .partitioning_kwargs --------------------
|
||||
|
||||
def it_provides_access_to_pass_through_kwargs_collected_by_the_partitioner_function(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["kwargs"] = {"foo": 42, "bar": "baz"}
|
||||
opts = MsgPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.partitioning_kwargs == {"foo": 42, "bar": "baz"}
|
||||
|
||||
# -- fixtures --------------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture
|
||||
def _last_modified_prop_(self, request: FixtureRequest):
|
||||
return property_mock(request, MsgPartitionerOptions, "_last_modified")
|
||||
|
||||
@pytest.fixture
|
||||
def Message_sent_date_(self, request: FixtureRequest):
|
||||
return property_mock(request, Message, "sent_date")
|
||||
|
||||
@pytest.fixture
|
||||
def opts_args(self) -> dict[str, Any]:
|
||||
"""All default arguments for `MsgPartitionerOptions`.
|
||||
|
||||
Individual argument values can be changed to suit each test. Makes construction of opts more
|
||||
compact for testing purposes.
|
||||
"""
|
||||
return {
|
||||
"file": None,
|
||||
"file_path": None,
|
||||
"metadata_file_path": None,
|
||||
"metadata_last_modified": None,
|
||||
"partition_attachments": False,
|
||||
"kwargs": {},
|
||||
}
|
||||
@@ -0,0 +1,311 @@
|
||||
"""Test-suite for `unstructured.partition.ndjson` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.unit_utils import example_doc_path
|
||||
from unstructured.documents.elements import CompositeElement
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.email import partition_email
|
||||
from unstructured.partition.html import partition_html
|
||||
from unstructured.partition.ndjson import partition_ndjson
|
||||
from unstructured.partition.text import partition_text
|
||||
from unstructured.partition.xml import partition_xml
|
||||
from unstructured.staging.base import elements_to_ndjson
|
||||
|
||||
DIRECTORY = pathlib.Path(__file__).parent.resolve()
|
||||
|
||||
is_in_docker = os.path.exists("/.dockerenv")
|
||||
|
||||
test_files = [
|
||||
"fake-text.txt",
|
||||
"fake-html.html",
|
||||
"eml/fake-email.eml",
|
||||
]
|
||||
|
||||
is_in_docker = os.path.exists("/.dockerenv")
|
||||
|
||||
|
||||
def test_it_chunks_elements_when_a_chunking_strategy_is_specified():
|
||||
chunks = partition_ndjson(
|
||||
example_doc_path("spring-weather.html.ndjson"),
|
||||
chunking_strategy="basic",
|
||||
max_characters=1500,
|
||||
)
|
||||
|
||||
assert len(chunks) == 9
|
||||
assert all(isinstance(ch, CompositeElement) for ch in chunks)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", test_files)
|
||||
def test_partition_ndjson_from_filename(filename: str):
|
||||
path = example_doc_path(filename)
|
||||
elements = []
|
||||
filetype = FileType.from_extension(os.path.splitext(path)[1])
|
||||
if filetype == FileType.TXT:
|
||||
elements = partition_text(filename=path)
|
||||
if filetype == FileType.HTML:
|
||||
elements = partition_html(filename=path)
|
||||
if filetype == FileType.XML:
|
||||
elements = partition_xml(filename=path)
|
||||
if filetype == FileType.EML:
|
||||
elements = partition_email(filename=path)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
_filename = os.path.basename(filename)
|
||||
test_path = os.path.join(tmpdir, _filename + ".ndjson")
|
||||
elements_to_ndjson(elements, filename=test_path)
|
||||
test_elements = partition_ndjson(filename=test_path)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert len(str(elements[0])) > 0
|
||||
|
||||
assert len(elements) == len(test_elements)
|
||||
for i in range(len(elements)):
|
||||
assert elements[i] == test_elements[i]
|
||||
assert elements[i].metadata.filename == filename.split("/")[-1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", test_files)
|
||||
def test_partition_ndjson_from_filename_with_metadata_filename(filename: str):
|
||||
path = example_doc_path(filename)
|
||||
elements = []
|
||||
filetype = FileType.from_extension(os.path.splitext(path)[1])
|
||||
if filetype == FileType.TXT:
|
||||
elements = partition_text(filename=path)
|
||||
if filetype == FileType.HTML:
|
||||
elements = partition_html(filename=path)
|
||||
if filetype == FileType.XML:
|
||||
elements = partition_xml(filename=path)
|
||||
if filetype == FileType.EML:
|
||||
elements = partition_email(filename=path)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
_filename = os.path.basename(filename)
|
||||
test_path = os.path.join(tmpdir, _filename + ".ndjson")
|
||||
elements_to_ndjson(elements, filename=test_path)
|
||||
test_elements = partition_ndjson(filename=test_path, metadata_filename="test")
|
||||
|
||||
assert len(test_elements) > 0
|
||||
assert len(str(test_elements[0])) > 0
|
||||
assert all(element.metadata.filename == "test" for element in test_elements)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", test_files)
|
||||
def test_partition_ndjson_from_file(filename: str):
|
||||
path = example_doc_path(filename)
|
||||
elements = []
|
||||
filetype = FileType.from_extension(os.path.splitext(path)[1])
|
||||
if filetype == FileType.TXT:
|
||||
elements = partition_text(filename=path)
|
||||
if filetype == FileType.HTML:
|
||||
elements = partition_html(filename=path)
|
||||
if filetype == FileType.XML:
|
||||
elements = partition_xml(filename=path)
|
||||
if filetype == FileType.EML:
|
||||
elements = partition_email(filename=path)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
_filename = os.path.basename(filename)
|
||||
test_path = os.path.join(tmpdir, _filename + ".ndjson")
|
||||
elements_to_ndjson(elements, filename=test_path)
|
||||
with open(test_path, "rb") as f:
|
||||
test_elements = partition_ndjson(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert len(str(elements[0])) > 0
|
||||
assert len(elements) == len(test_elements)
|
||||
for i in range(len(elements)):
|
||||
assert elements[i] == test_elements[i]
|
||||
assert elements[i].metadata.filename == filename.split("/")[-1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", test_files)
|
||||
def test_partition_ndjson_from_file_with_metadata_filename(filename: str):
|
||||
path = example_doc_path(filename)
|
||||
elements = []
|
||||
filetype = FileType.from_extension(os.path.splitext(path)[1])
|
||||
if filetype == FileType.TXT:
|
||||
elements = partition_text(filename=path)
|
||||
if filetype == FileType.HTML:
|
||||
elements = partition_html(filename=path)
|
||||
if filetype == FileType.XML:
|
||||
elements = partition_xml(filename=path)
|
||||
if filetype == FileType.EML:
|
||||
elements = partition_email(filename=path)
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
_filename = os.path.basename(filename)
|
||||
test_path = os.path.join(tmpdir, _filename + ".ndjson")
|
||||
elements_to_ndjson(elements, filename=test_path)
|
||||
with open(test_path, "rb") as f:
|
||||
test_elements = partition_ndjson(file=f, metadata_filename="test")
|
||||
|
||||
for i in range(len(test_elements)):
|
||||
assert test_elements[i].metadata.filename == "test"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", test_files)
|
||||
def test_partition_ndjson_from_text(filename: str):
|
||||
path = example_doc_path(filename)
|
||||
elements = []
|
||||
filetype = FileType.from_extension(os.path.splitext(path)[1])
|
||||
if filetype == FileType.TXT:
|
||||
elements = partition_text(filename=path)
|
||||
if filetype == FileType.HTML:
|
||||
elements = partition_html(filename=path)
|
||||
if filetype == FileType.XML:
|
||||
elements = partition_xml(filename=path)
|
||||
if filetype == FileType.EML:
|
||||
elements = partition_email(filename=path)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
_filename = os.path.basename(filename)
|
||||
test_path = os.path.join(tmpdir, _filename + ".ndjson")
|
||||
elements_to_ndjson(elements, filename=test_path)
|
||||
with open(test_path) as f:
|
||||
text = f.read()
|
||||
test_elements = partition_ndjson(text=text)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert len(str(elements[0])) > 0
|
||||
assert len(elements) == len(test_elements)
|
||||
for i in range(len(elements)):
|
||||
assert elements[i] == test_elements[i]
|
||||
assert elements[i].metadata.filename == filename.split("/")[-1]
|
||||
|
||||
|
||||
def test_partition_json_raises_with_none_specified():
|
||||
with pytest.raises(ValueError):
|
||||
partition_ndjson()
|
||||
|
||||
|
||||
def test_partition_ndjson_works_with_empty_string():
|
||||
assert partition_ndjson(text="") == []
|
||||
|
||||
|
||||
def test_partition_ndjson_fails_with_empty_item():
|
||||
with pytest.raises(ValueError):
|
||||
partition_ndjson(text="{}")
|
||||
|
||||
|
||||
def test_partition_ndjson_fails_with_empty_list():
|
||||
with pytest.raises(ValueError):
|
||||
partition_ndjson(text="[]")
|
||||
|
||||
|
||||
def test_partition_ndjson_raises_with_too_many_specified():
|
||||
path = example_doc_path("fake-text.txt")
|
||||
elements = []
|
||||
filetype = FileType.from_extension(os.path.splitext(path)[1])
|
||||
if filetype == FileType.TXT:
|
||||
elements = partition_text(filename=path)
|
||||
if filetype == FileType.HTML:
|
||||
elements = partition_html(filename=path)
|
||||
if filetype == FileType.XML:
|
||||
elements = partition_xml(filename=path)
|
||||
if filetype == FileType.EML:
|
||||
elements = partition_email(filename=path)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
test_path = os.path.join(tmpdir, "fake-text.txt.ndjson")
|
||||
elements_to_ndjson(elements, filename=test_path)
|
||||
with open(test_path, "rb") as f:
|
||||
text = f.read().decode("utf-8")
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_ndjson(filename=test_path, file=f)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_ndjson(filename=test_path, text=text)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_ndjson(file=f, text=text)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_ndjson(filename=test_path, file=f, text=text)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_ndjson_from_file_path_gets_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.ndjson.get_last_modified_date",
|
||||
return_value=filesystem_last_modified,
|
||||
)
|
||||
|
||||
elements = partition_ndjson(example_doc_path("spring-weather.html.ndjson"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_ndjson_from_file_gets_last_modified_None():
|
||||
with open(example_doc_path("spring-weather.html.ndjson"), "rb") as f:
|
||||
elements = partition_ndjson(file=f)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_ndjson_from_text_gets_last_modified_None():
|
||||
with open(example_doc_path("spring-weather.html.ndjson")) as f:
|
||||
text = f.read()
|
||||
|
||||
elements = partition_ndjson(text=text)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_ndjson_from_file_path_prefers_metadata_last_modified(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.ndjson.get_last_modified_date",
|
||||
return_value=filesystem_last_modified,
|
||||
)
|
||||
|
||||
elements = partition_ndjson(
|
||||
example_doc_path("spring-weather.html.ndjson"),
|
||||
metadata_last_modified=metadata_last_modified,
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_ndjson_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
with open(example_doc_path("spring-weather.html.ndjson"), "rb") as f:
|
||||
elements = partition_ndjson(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_ndjson_from_text_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
with open(example_doc_path("spring-weather.html.ndjson")) as f:
|
||||
text = f.read()
|
||||
|
||||
elements = partition_ndjson(text=text, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_json_raises_with_unprocessable_json():
|
||||
text = '{"invalid": "schema"}'
|
||||
with pytest.raises(ValueError):
|
||||
partition_ndjson(text=text)
|
||||
|
||||
|
||||
def test_partition_json_raises_with_invalid_json():
|
||||
text = '[{"hi": "there"}]]'
|
||||
with pytest.raises(ValueError):
|
||||
partition_ndjson(text=text)
|
||||
@@ -0,0 +1,224 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
"""Test suite for `unstructured.partition.odt` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Iterator
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.unit_utils import (
|
||||
ANY,
|
||||
FixtureRequest,
|
||||
assert_round_trips_through_JSON,
|
||||
example_doc_path,
|
||||
method_mock,
|
||||
)
|
||||
from unstructured.chunking.basic import chunk_elements
|
||||
from unstructured.documents.elements import (
|
||||
CompositeElement,
|
||||
Element,
|
||||
Table,
|
||||
TableChunk,
|
||||
Text,
|
||||
)
|
||||
from unstructured.partition.docx import partition_docx
|
||||
from unstructured.partition.odt import partition_odt
|
||||
from unstructured.partition.utils.constants import UNSTRUCTURED_INCLUDE_DEBUG_METADATA
|
||||
|
||||
|
||||
def test_partition_odt_matches_partition_docx():
|
||||
odt_file_path = example_doc_path("simple.odt")
|
||||
docx_file_path = example_doc_path("simple.docx")
|
||||
|
||||
assert partition_odt(odt_file_path) == partition_docx(docx_file_path)
|
||||
|
||||
|
||||
# -- document-source (file or filename) ----------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_odt_from_filename():
|
||||
elements = partition_odt(example_doc_path("fake.odt"))
|
||||
|
||||
assert elements == [
|
||||
Text("Lorem ipsum dolor sit amet."),
|
||||
Table(
|
||||
"Header row Mon Wed Fri"
|
||||
" Color Blue Red Green"
|
||||
" Time 1pm 2pm 3pm"
|
||||
" Leader Sarah Mark Ryan"
|
||||
),
|
||||
]
|
||||
assert all(e.metadata.filename == "fake.odt" for e in elements)
|
||||
if UNSTRUCTURED_INCLUDE_DEBUG_METADATA:
|
||||
# -- document is ultimately partitioned by partition_docx() --
|
||||
assert {e.metadata.detection_origin for e in elements} == {"docx"}
|
||||
|
||||
|
||||
def test_partition_odt_from_file():
|
||||
with open(example_doc_path("fake.odt"), "rb") as f:
|
||||
elements = partition_odt(file=f)
|
||||
|
||||
assert elements == [
|
||||
Text("Lorem ipsum dolor sit amet."),
|
||||
Table(
|
||||
"Header row Mon Wed Fri"
|
||||
" Color Blue Red Green"
|
||||
" Time 1pm 2pm 3pm"
|
||||
" Leader Sarah Mark Ryan"
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
# -- .metadata.filename --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_odt_from_filename_gets_the_ODT_filename_in_metadata_not_the_DOCX_filename():
|
||||
elements = partition_odt(example_doc_path("simple.odt"))
|
||||
assert all(e.metadata.filename == "simple.odt" for e in elements), (
|
||||
f"Expected all elements to have 'simple.odt' as their filename, but got:"
|
||||
f" {repr(elements[0].metadata.filename)}"
|
||||
)
|
||||
|
||||
|
||||
def test_partition_odt_from_filename_with_metadata_filename():
|
||||
elements = partition_odt(example_doc_path("fake.odt"), metadata_filename="test")
|
||||
assert all(e.metadata.filename == "test" for e in elements)
|
||||
|
||||
|
||||
def test_partition_odt_from_file_with_metadata_filename():
|
||||
with open(example_doc_path("fake.odt"), "rb") as f:
|
||||
elements = partition_odt(file=f, metadata_filename="test")
|
||||
assert all(e.metadata.filename == "test" for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.filetype --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_odt_gets_the_ODT_MIME_type_in_metadata_filetype():
|
||||
ODT_MIME_TYPE = "application/vnd.oasis.opendocument.text"
|
||||
elements = partition_odt(example_doc_path("simple.odt"))
|
||||
assert all(e.metadata.filetype == ODT_MIME_TYPE for e in elements), (
|
||||
f"Expected all elements to have '{ODT_MIME_TYPE}' as their filetype, but got:"
|
||||
f" {repr(elements[0].metadata.filetype)}"
|
||||
)
|
||||
|
||||
|
||||
# -- .metadata.text_as_html ----------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kwargs", [{}, {"infer_table_structure": True}])
|
||||
def test_partition_odt_adds_text_as_html_when_infer_table_structure_is_omitted_or_True(
|
||||
kwargs: dict[str, Any],
|
||||
):
|
||||
with open(example_doc_path("fake.odt"), "rb") as f:
|
||||
elements = partition_odt(file=f, **kwargs)
|
||||
|
||||
table = elements[1]
|
||||
assert isinstance(table, Table)
|
||||
assert table.metadata.text_as_html is not None
|
||||
assert table.metadata.text_as_html.startswith("<table>")
|
||||
|
||||
|
||||
def test_partition_odt_suppresses_text_as_html_when_infer_table_structure_is_False():
|
||||
with open(example_doc_path("fake.odt"), "rb") as f:
|
||||
elements = partition_odt(file=f, infer_table_structure=False)
|
||||
|
||||
table = elements[1]
|
||||
assert isinstance(table, Table)
|
||||
assert table.metadata.text_as_html is None
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_odt_pulls_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.odt.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_odt(example_doc_path("fake.odt"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_odt_prefers_metadata_last_modified_when_provided(mocker: MockFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.odt.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_odt(
|
||||
example_doc_path("simple.odt"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.languages -------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_odt_adds_languages_metadata():
|
||||
elements = partition_odt(example_doc_path("simple.odt"))
|
||||
assert all(e.metadata.languages == ["eng"] for e in elements)
|
||||
|
||||
|
||||
def test_partition_odt_respects_detect_language_per_element_arg():
|
||||
elements = partition_odt(
|
||||
example_doc_path("language-docs/eng_spa_mult.odt"), detect_language_per_element=True
|
||||
)
|
||||
assert [e.metadata.languages for e in elements] == [
|
||||
["eng"],
|
||||
["spa", "eng"],
|
||||
["eng"],
|
||||
["eng"],
|
||||
["spa"],
|
||||
]
|
||||
|
||||
|
||||
# -- miscellaneous -------------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "expected_value"),
|
||||
[({}, "hi_res"), ({"strategy": None}, "hi_res"), ({"strategy": "auto"}, "auto")],
|
||||
)
|
||||
def test_partition_odt_forwards_strategy_arg_to_partition_docx(
|
||||
request: FixtureRequest, kwargs: dict[str, Any], expected_value: str | None
|
||||
):
|
||||
from unstructured.partition.docx import _DocxPartitioner
|
||||
|
||||
def fake_iter_document_elements(self: _DocxPartitioner) -> Iterator[Element]:
|
||||
yield Text(f"strategy == {self._opts.strategy}")
|
||||
|
||||
_iter_elements_ = method_mock(
|
||||
request,
|
||||
_DocxPartitioner,
|
||||
"_iter_document_elements",
|
||||
side_effect=fake_iter_document_elements,
|
||||
)
|
||||
|
||||
(element,) = partition_odt(example_doc_path("simple.odt"), **kwargs)
|
||||
|
||||
_iter_elements_.assert_called_once_with(ANY)
|
||||
assert element.text == f"strategy == {expected_value}"
|
||||
|
||||
|
||||
def test_partition_odt_round_trips_through_json():
|
||||
"""Elements produced can be serialized then deserialized without loss."""
|
||||
assert_round_trips_through_JSON(partition_odt(example_doc_path("simple.odt")))
|
||||
|
||||
|
||||
def test_partition_odt_chunks_elements_when_chunking_strategy_is_specified():
|
||||
document_path = example_doc_path("simple.odt")
|
||||
elements = partition_odt(document_path)
|
||||
chunks = partition_odt(document_path, chunking_strategy="basic")
|
||||
|
||||
# -- all chunks are chunk element-types --
|
||||
assert all(isinstance(c, (CompositeElement, Table, TableChunk)) for c in chunks)
|
||||
# -- chunks from partitioning match those produced by chunking elements in separate step --
|
||||
assert chunks == chunk_elements(elements)
|
||||
@@ -0,0 +1,158 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.unit_utils import (
|
||||
assert_round_trips_through_JSON,
|
||||
example_doc_path,
|
||||
find_text_in_elements,
|
||||
)
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import Title
|
||||
from unstructured.partition.org import partition_org
|
||||
|
||||
|
||||
def test_partition_org_from_filename():
|
||||
elements = partition_org(example_doc_path("README.org"))
|
||||
|
||||
assert elements[0] == Title("Example Docs")
|
||||
assert elements[0].metadata.filetype == "text/org"
|
||||
|
||||
|
||||
def test_partition_org_from_file():
|
||||
with open(example_doc_path("README.org"), "rb") as f:
|
||||
elements = partition_org(file=f)
|
||||
|
||||
assert elements[0] == Title("Example Docs")
|
||||
|
||||
|
||||
# -- .metadata.filename --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_org_from_filename_gets_filename_from_filename_arg():
|
||||
elements = partition_org(example_doc_path("README.org"))
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename == "README.org" for e in elements)
|
||||
|
||||
|
||||
def test_partition_org_from_file_gets_filename_None():
|
||||
with open(example_doc_path("README.org"), "rb") as f:
|
||||
elements = partition_org(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_org_from_filename_prefers_metadata_filename():
|
||||
elements = partition_org(example_doc_path("README.org"), metadata_filename="orig-name.org")
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(element.metadata.filename == "orig-name.org" for element in elements)
|
||||
|
||||
|
||||
def test_partition_org_from_file_prefers_metadata_filename():
|
||||
with open(example_doc_path("README.org"), "rb") as f:
|
||||
elements = partition_org(file=f, metadata_filename="orig-name.org")
|
||||
|
||||
assert all(e.metadata.filename == "orig-name.org" for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.filetype --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_org_gets_the_ORG_MIME_type_in_metadata_filetype():
|
||||
ORG_MIME_TYPE = "text/org"
|
||||
elements = partition_org(example_doc_path("README.org"))
|
||||
assert all(e.metadata.filetype == ORG_MIME_TYPE for e in elements), (
|
||||
f"Expected all elements to have '{ORG_MIME_TYPE}' as their filetype, but got:"
|
||||
f" {repr(elements[0].metadata.filetype)}"
|
||||
)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_org_from_filename_gets_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2024-06-14T16:01:29"
|
||||
mocker.patch(
|
||||
"unstructured.partition.org.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_org(example_doc_path("README.org"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_org_from_file_gets_last_modified_None():
|
||||
with open(example_doc_path("README.org"), "rb") as f:
|
||||
elements = partition_org(file=f)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_org_from_filename_prefers_metadata_last_modified(mocker: MockFixture):
|
||||
filesystem_last_modified = "2020-08-04T06:11:47"
|
||||
metadata_last_modified = "2024-06-14T16:01:29"
|
||||
mocker.patch(
|
||||
"unstructured.partition.org.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_org(
|
||||
example_doc_path("README.org"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_org_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
with open(example_doc_path("README.org"), "rb") as f:
|
||||
elements = partition_org(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_org_with_json():
|
||||
elements = partition_org(example_doc_path("README.org"))
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
def test_add_chunking_strategy_by_title_on_partition_org():
|
||||
file_path = example_doc_path("README.org")
|
||||
elements = partition_org(file_path)
|
||||
chunk_elements = partition_org(file_path, chunking_strategy="by_title")
|
||||
chunks = chunk_by_title(elements)
|
||||
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_partition_org_element_metadata_has_languages():
|
||||
elements = partition_org(example_doc_path("README.org"))
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_org_respects_detect_language_per_element():
|
||||
elements = partition_org(
|
||||
example_doc_path("language-docs/eng_spa_mult.org"), detect_language_per_element=True
|
||||
)
|
||||
langs = [element.metadata.languages for element in elements]
|
||||
assert langs == [["eng"], ["spa", "eng"], ["eng"], ["eng"], ["spa"]]
|
||||
|
||||
|
||||
def test_org_wont_include_external_files():
|
||||
# Make sure our import file is in place (otherwise the import fails silently and test passes)
|
||||
assert Path(example_doc_path("file_we_dont_want_imported")).exists()
|
||||
elements = partition_org(example_doc_path("README-w-include.org"))
|
||||
# The partition should contain some elements
|
||||
assert elements
|
||||
# We find something we expect to find from file we partitioned directly
|
||||
assert find_text_in_elements("instructions", elements)
|
||||
# But we don't find something from the file included within the file we partitioned directly
|
||||
assert not find_text_in_elements("wombat", elements)
|
||||
@@ -0,0 +1,180 @@
|
||||
"""Test-suite for `unstructured.partition.ppt` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.unit_utils import assert_round_trips_through_JSON, example_doc_path
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import ListItem, NarrativeText, PageBreak, Title
|
||||
from unstructured.partition.ppt import partition_ppt
|
||||
from unstructured.partition.utils.constants import UNSTRUCTURED_INCLUDE_DEBUG_METADATA
|
||||
|
||||
EXPECTED_PPT_OUTPUT = [
|
||||
Title(text="Adding a Bullet Slide"),
|
||||
ListItem(text="Find the bullet slide layout"),
|
||||
ListItem(text="Use _TextFrame.text for first bullet"),
|
||||
ListItem(text="Use _TextFrame.add_paragraph() for subsequent bullets"),
|
||||
NarrativeText(text="Here is a lot of text!"),
|
||||
NarrativeText(text="Here is some text in a text box!"),
|
||||
]
|
||||
|
||||
|
||||
def test_partition_ppt_from_filename():
|
||||
elements = partition_ppt(example_doc_path("fake-power-point.ppt"))
|
||||
assert elements == EXPECTED_PPT_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "fake-power-point.ppt"
|
||||
if UNSTRUCTURED_INCLUDE_DEBUG_METADATA:
|
||||
assert {element.metadata.detection_origin for element in elements} == {"pptx"}
|
||||
|
||||
|
||||
def test_partition_ppt_raises_with_missing_file():
|
||||
with pytest.raises(ValueError):
|
||||
partition_ppt(example_doc_path("doesnt-exist.ppt"))
|
||||
|
||||
|
||||
def test_partition_ppt_from_file():
|
||||
with open(example_doc_path("fake-power-point.ppt"), "rb") as f:
|
||||
elements = partition_ppt(file=f)
|
||||
assert elements == EXPECTED_PPT_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename is None
|
||||
|
||||
|
||||
def test_partition_ppt_from_file_with_metadata_filename():
|
||||
with open(example_doc_path("fake-power-point.ppt"), "rb") as f:
|
||||
elements = partition_ppt(file=f, metadata_filename="test")
|
||||
assert elements == EXPECTED_PPT_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "test"
|
||||
|
||||
|
||||
def test_partition_ppt_raises_with_both_specified():
|
||||
filename = example_doc_path("fake-power-point.ppt")
|
||||
with open(filename, "rb") as f, pytest.raises(ValueError):
|
||||
partition_ppt(filename=filename, file=f)
|
||||
|
||||
|
||||
def test_partition_ppt_raises_when_neither_file_path_or_file_is_provided():
|
||||
with pytest.raises(ValueError):
|
||||
partition_ppt()
|
||||
|
||||
|
||||
# -- .metadata.filename --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_ppt_from_filename_gets_filename_from_filename_arg():
|
||||
elements = partition_ppt(example_doc_path("fake-power-point.ppt"))
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename == "fake-power-point.ppt" for e in elements)
|
||||
|
||||
|
||||
def test_partition_ppt_from_file_gets_filename_None():
|
||||
with open(example_doc_path("fake-power-point.ppt"), "rb") as f:
|
||||
elements = partition_ppt(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_ppt_from_filename_prefers_metadata_filename():
|
||||
elements = partition_ppt(example_doc_path("fake-power-point.ppt"), metadata_filename="test")
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(element.metadata.filename == "test" for element in elements)
|
||||
|
||||
|
||||
def test_partition_ppt_from_file_prefers_metadata_filename():
|
||||
with open(example_doc_path("fake-power-point.ppt"), "rb") as f:
|
||||
elements = partition_ppt(file=f, metadata_filename="test")
|
||||
|
||||
assert all(e.metadata.filename == "test" for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_ppt_from_file_path_gets_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2024-05-01T15:37:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.ppt.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_ppt(example_doc_path("fake-power-point.ppt"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_ppt_from_file_gets_last_modified_None():
|
||||
with open(example_doc_path("fake-power-point.ppt"), "rb") as f:
|
||||
elements = partition_ppt(file=f)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_ppt_from_file_path_prefers_metadata_last_modified(mocker: MockFixture):
|
||||
filesystem_last_modified = "2024-05-01T15:37:28"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.ppt.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_ppt(
|
||||
example_doc_path("fake-power-point.ppt"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_ppt_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
|
||||
with open(example_doc_path("fake-power-point.ppt"), "rb") as f:
|
||||
elements = partition_ppt(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert elements[0].metadata.last_modified == metadata_last_modified
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_ppt_with_json():
|
||||
elements = partition_ppt(example_doc_path("fake-power-point.ppt"))
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
def test_add_chunking_strategy_by_title_on_partition_ppt():
|
||||
file_path = example_doc_path("fake-power-point.ppt")
|
||||
elements = partition_ppt(file_path)
|
||||
chunk_elements = partition_ppt(file_path, chunking_strategy="by_title")
|
||||
chunks = chunk_by_title(elements)
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_partition_ppt_params():
|
||||
"""Integration test of params: languages, include_page_break, and include_slide_notes."""
|
||||
elements = partition_ppt(
|
||||
example_doc_path("language-docs/eng_spa_mult.ppt"),
|
||||
include_page_breaks=True,
|
||||
include_slide_notes=True,
|
||||
)
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
assert any(isinstance(element, PageBreak) for element in elements)
|
||||
# The example doc contains a slide note with the text "This is a slide note."
|
||||
assert any(element.text == "This is a slide note." for element in elements)
|
||||
|
||||
|
||||
def test_partition_ppt_respects_detect_language_per_element():
|
||||
elements = partition_ppt(
|
||||
example_doc_path("language-docs/eng_spa_mult.ppt"), detect_language_per_element=True
|
||||
)
|
||||
langs = [element.metadata.languages for element in elements]
|
||||
# languages other than English and Spanish are detected by this partitioner,
|
||||
# so this test is slightly different from the other partition tests
|
||||
langs = {element.metadata.languages[0] for element in elements if element.metadata.languages}
|
||||
assert "eng" in langs
|
||||
assert "spa" in langs
|
||||
@@ -0,0 +1,799 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
"""Test suite for `unstructured.partition.pptx` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import io
|
||||
import pathlib
|
||||
import tempfile
|
||||
from typing import Any, Iterator, cast
|
||||
|
||||
import pptx
|
||||
import pytest
|
||||
from pptx.shapes.picture import Picture
|
||||
from pptx.util import Inches
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.unit_utils import (
|
||||
FixtureRequest,
|
||||
Mock,
|
||||
assert_round_trips_through_JSON,
|
||||
example_doc_path,
|
||||
function_mock,
|
||||
property_mock,
|
||||
)
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import (
|
||||
Element,
|
||||
ElementMetadata,
|
||||
Image,
|
||||
ListItem,
|
||||
NarrativeText,
|
||||
PageBreak,
|
||||
Text,
|
||||
Title,
|
||||
)
|
||||
from unstructured.partition.pptx import (
|
||||
PptxPartitionerOptions,
|
||||
_PptxPartitioner,
|
||||
partition_pptx,
|
||||
register_picture_partitioner,
|
||||
)
|
||||
|
||||
EXPECTED_PPTX_OUTPUT = [
|
||||
Title(text="Adding a Bullet Slide"),
|
||||
ListItem(text="Find the bullet slide layout"),
|
||||
ListItem(text="Use _TextFrame.text for first bullet"),
|
||||
ListItem(text="Use _TextFrame.add_paragraph() for subsequent bullets"),
|
||||
NarrativeText(text="Here is a lot of text!"),
|
||||
NarrativeText(text="Here is some text in a text box!"),
|
||||
]
|
||||
|
||||
|
||||
# == document file behaviors =====================================================================
|
||||
|
||||
|
||||
def test_partition_pptx_from_filename():
|
||||
elements = partition_pptx(example_doc_path("fake-power-point.pptx"))
|
||||
assert elements == EXPECTED_PPTX_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "fake-power-point.pptx"
|
||||
|
||||
|
||||
def test_partition_pptx_from_filename_with_metadata_filename():
|
||||
elements = partition_pptx(example_doc_path("fake-power-point.pptx"), metadata_filename="test")
|
||||
assert elements == EXPECTED_PPTX_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "test"
|
||||
|
||||
|
||||
def test_partition_pptx_with_spooled_file():
|
||||
"""The `partition_pptx() function can handle a `SpooledTemporaryFile.
|
||||
|
||||
Including one that does not have its read-pointer set to the start.
|
||||
"""
|
||||
with tempfile.SpooledTemporaryFile() as spooled_temp_file:
|
||||
with open(example_doc_path("fake-power-point.pptx"), "rb") as test_file:
|
||||
spooled_temp_file.write(test_file.read())
|
||||
|
||||
elements = partition_pptx(file=spooled_temp_file)
|
||||
|
||||
assert elements == EXPECTED_PPTX_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename is None
|
||||
|
||||
|
||||
def test_partition_pptx_from_file():
|
||||
with open(example_doc_path("fake-power-point.pptx"), "rb") as f:
|
||||
elements = partition_pptx(file=f)
|
||||
assert elements == EXPECTED_PPTX_OUTPUT
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_pptx_from_file_with_metadata_filename():
|
||||
with open(example_doc_path("fake-power-point.pptx"), "rb") as f:
|
||||
elements = partition_pptx(file=f, metadata_filename="test")
|
||||
assert elements == EXPECTED_PPTX_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "test"
|
||||
|
||||
|
||||
def test_partition_pptx_raises_with_neither():
|
||||
with pytest.raises(ValueError):
|
||||
partition_pptx()
|
||||
|
||||
|
||||
def test_partition_pptx_recurses_into_group_shapes():
|
||||
elements = partition_pptx(example_doc_path("group-shapes-nested.pptx"))
|
||||
assert [e.text for e in elements] == ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"]
|
||||
|
||||
|
||||
def test_it_loads_a_PPTX_with_a_JPEG_misidentified_as_image_jpg(opts_args: dict[str, Any]):
|
||||
opts_args["file_path"] = example_doc_path("test-image-jpg-mime.pptx")
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
prs = _PptxPartitioner(opts)._presentation
|
||||
picture = cast(Picture, prs.slides[0].shapes[0])
|
||||
|
||||
try:
|
||||
picture.image
|
||||
except AttributeError:
|
||||
raise AssertionError("JPEG image not recognized, needs `python-pptx>=1.0.1`")
|
||||
|
||||
|
||||
# == page-break behaviors ========================================================================
|
||||
|
||||
|
||||
def test_partition_pptx_adds_page_breaks(tmp_path: pathlib.Path):
|
||||
filename = str(tmp_path / "test-page-breaks.pptx")
|
||||
|
||||
presentation = pptx.Presentation()
|
||||
blank_slide_layout = presentation.slide_layouts[6]
|
||||
|
||||
slide = presentation.slides.add_slide(blank_slide_layout)
|
||||
left = top = width = height = Inches(2)
|
||||
txBox = slide.shapes.add_textbox(left, top, width, height)
|
||||
tf = txBox.text_frame
|
||||
tf.text = "This is the first slide."
|
||||
|
||||
slide = presentation.slides.add_slide(blank_slide_layout)
|
||||
left = top = width = height = Inches(2)
|
||||
txBox = slide.shapes.add_textbox(left, top, width, height)
|
||||
tf = txBox.text_frame
|
||||
tf.text = "This is the second slide."
|
||||
|
||||
presentation.save(filename)
|
||||
|
||||
elements = partition_pptx(filename=filename)
|
||||
|
||||
assert elements == [
|
||||
NarrativeText(text="This is the first slide."),
|
||||
PageBreak(text=""),
|
||||
NarrativeText(text="This is the second slide."),
|
||||
]
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "test-page-breaks.pptx"
|
||||
|
||||
|
||||
def test_partition_pptx_page_breaks_toggle_off(tmp_path: pathlib.Path):
|
||||
filename = str(tmp_path / "test-page-breaks.pptx")
|
||||
|
||||
presentation = pptx.Presentation()
|
||||
blank_slide_layout = presentation.slide_layouts[6]
|
||||
|
||||
slide = presentation.slides.add_slide(blank_slide_layout)
|
||||
left = top = width = height = Inches(2)
|
||||
txBox = slide.shapes.add_textbox(left, top, width, height)
|
||||
tf = txBox.text_frame
|
||||
tf.text = "This is the first slide."
|
||||
|
||||
slide = presentation.slides.add_slide(blank_slide_layout)
|
||||
left = top = width = height = Inches(2)
|
||||
txBox = slide.shapes.add_textbox(left, top, width, height)
|
||||
tf = txBox.text_frame
|
||||
tf.text = "This is the second slide."
|
||||
|
||||
presentation.save(filename)
|
||||
|
||||
elements = partition_pptx(filename=filename, include_page_breaks=False)
|
||||
|
||||
assert elements == [
|
||||
NarrativeText(text="This is the first slide."),
|
||||
NarrativeText(text="This is the second slide."),
|
||||
]
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "test-page-breaks.pptx"
|
||||
|
||||
|
||||
def test_partition_pptx_many_pages():
|
||||
elements = partition_pptx(example_doc_path("fake-power-point-many-pages.pptx"))
|
||||
|
||||
# The page_number of PageBreak is None
|
||||
assert set(filter(None, (elt.metadata.page_number for elt in elements))) == {1, 2}
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "fake-power-point-many-pages.pptx"
|
||||
|
||||
|
||||
# == miscellaneous behaviors =====================================================================
|
||||
|
||||
|
||||
def test_partition_pptx_orders_elements(tmp_path: pathlib.Path):
|
||||
filename = str(tmp_path / "test-ordering.pptx")
|
||||
presentation = pptx.Presentation()
|
||||
blank_slide_layout = presentation.slide_layouts[6]
|
||||
slide = presentation.slides.add_slide(blank_slide_layout)
|
||||
|
||||
left = top = width = height = Inches(2)
|
||||
txBox = slide.shapes.add_textbox(left, top, width, height)
|
||||
tf = txBox.text_frame
|
||||
tf.text = "This is lower and should come second"
|
||||
|
||||
left = top = width = height = Inches(1)
|
||||
left = top = Inches(-10)
|
||||
txBox = slide.shapes.add_textbox(left, top, width, height)
|
||||
tf = txBox.text_frame
|
||||
tf.text = "This is off the page and shouldn't appear"
|
||||
|
||||
left = top = width = height = Inches(2)
|
||||
txBox = slide.shapes.add_textbox(left, top, width, height)
|
||||
tf = txBox.text_frame
|
||||
tf.text = ""
|
||||
|
||||
left = top = width = height = Inches(1)
|
||||
txBox = slide.shapes.add_textbox(left, top, width, height)
|
||||
tf = txBox.text_frame
|
||||
tf.text = "This is higher and should come first"
|
||||
|
||||
top = width = height = Inches(1)
|
||||
left = Inches(0.5)
|
||||
txBox = slide.shapes.add_textbox(left, top, width, height)
|
||||
tf = txBox.text_frame
|
||||
tf.text = "-------------TOP-------------"
|
||||
|
||||
presentation.save(filename)
|
||||
|
||||
elements = partition_pptx(filename=filename)
|
||||
assert elements == [
|
||||
Text("-------------TOP-------------"),
|
||||
NarrativeText("This is higher and should come first"),
|
||||
NarrativeText("This is lower and should come second"),
|
||||
]
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "test-ordering.pptx"
|
||||
|
||||
|
||||
def test_partition_pptx_grabs_tables():
|
||||
elements = partition_pptx(example_doc_path("fake-power-point-table.pptx"))
|
||||
|
||||
assert elements[1].text.startswith("Column 1")
|
||||
assert elements[1].text.strip().endswith("Aqua")
|
||||
assert elements[1].metadata.text_as_html == (
|
||||
"<table>"
|
||||
"<tr><td>Column 1</td><td>Column 2</td><td>Column 3</td></tr>"
|
||||
"<tr><td>Red</td><td>Green</td><td>Blue</td></tr>"
|
||||
"<tr><td>Purple</td><td>Orange</td><td>Yellow</td></tr>"
|
||||
"<tr><td>Tangerine</td><td>Pink</td><td>Aqua</td></tr>"
|
||||
"</table>"
|
||||
)
|
||||
assert elements[1].metadata.filename == "fake-power-point-table.pptx"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("infer_table_structure", [True, False])
|
||||
def test_partition_pptx_infer_table_structure(infer_table_structure: bool):
|
||||
elements = partition_pptx(
|
||||
example_doc_path("fake-power-point-table.pptx"), infer_table_structure=infer_table_structure
|
||||
)
|
||||
table_element_has_text_as_html_field = (
|
||||
hasattr(elements[1].metadata, "text_as_html")
|
||||
and elements[1].metadata.text_as_html is not None
|
||||
)
|
||||
assert table_element_has_text_as_html_field == infer_table_structure
|
||||
|
||||
|
||||
def test_partition_pptx_malformed():
|
||||
elements = partition_pptx(example_doc_path("fake-power-point-malformed.pptx"))
|
||||
|
||||
assert elements[0].text == "Problem Date Placeholder"
|
||||
assert elements[1].text == "Test Slide"
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "fake-power-point-malformed.pptx"
|
||||
|
||||
|
||||
# == image sub-partitioning behaviors ============================================================
|
||||
|
||||
|
||||
def test_partition_pptx_generates_no_Image_elements_by_default():
|
||||
assert partition_pptx(example_doc_path("picture.pptx")) == []
|
||||
|
||||
|
||||
def test_partition_pptx_uses_registered_picture_partitioner():
|
||||
class FakePicturePartitioner:
|
||||
@classmethod
|
||||
def iter_elements(cls, picture: Picture, opts: PptxPartitionerOptions) -> Iterator[Element]:
|
||||
image_hash = hashlib.sha1(picture.image.blob).hexdigest()
|
||||
yield Image(f"Image with hash {image_hash}, strategy: {opts.strategy}")
|
||||
|
||||
register_picture_partitioner(FakePicturePartitioner)
|
||||
|
||||
elements = partition_pptx(example_doc_path("picture.pptx"))
|
||||
|
||||
assert len(elements) == 1
|
||||
image = elements[0]
|
||||
assert type(image) is Image
|
||||
assert image.text == "Image with hash b0a1e6cf904691e6fa42bd9e72acc2b05280dc86, strategy: fast"
|
||||
|
||||
|
||||
# == metadata behaviors ==========================================================================
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_pptx_from_file_path_gets_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2024-05-01T15:37:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.pptx.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_pptx(example_doc_path("simple.pptx"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_pptx_from_file_gets_last_modified_None():
|
||||
with open(example_doc_path("simple.pptx"), "rb") as f:
|
||||
elements = partition_pptx(file=f)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_pptx_from_file_path_prefers_metadata_last_modified(mocker: MockFixture):
|
||||
filesystem_last_modified = "2024-05-01T15:37:28"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.pptx.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_pptx(
|
||||
example_doc_path("simple.pptx"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_pptx_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
with open(example_doc_path("simple.pptx"), "rb") as f:
|
||||
elements = partition_pptx(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.languages -------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_pptx_element_metadata_has_languages():
|
||||
elements = partition_pptx(example_doc_path("fake-power-point.pptx"))
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_pptx_respects_detect_language_per_element():
|
||||
elements = partition_pptx(
|
||||
example_doc_path("language-docs/eng_spa_mult.pptx"), detect_language_per_element=True
|
||||
)
|
||||
|
||||
langs = [element.metadata.languages for element in elements]
|
||||
# languages other than English and Spanish are detected by this partitioner,
|
||||
# so this test is slightly different from the other partition tests
|
||||
langs = {element.metadata.languages[0] for element in elements if element.metadata.languages}
|
||||
assert "eng" in langs
|
||||
assert "spa" in langs
|
||||
|
||||
|
||||
def test_partition_pptx_raises_TypeError_for_invalid_languages():
|
||||
with pytest.raises(TypeError):
|
||||
partition_pptx(example_doc_path("fake-power-point.pptx"), languages="eng")
|
||||
|
||||
|
||||
# == downstream behaviors ========================================================================
|
||||
|
||||
|
||||
def test_partition_pptx_with_json():
|
||||
elements = partition_pptx(example_doc_path("fake-power-point.pptx"))
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
def test_add_chunking_strategy_by_title_on_partition_pptx():
|
||||
filename = example_doc_path("science-exploration-1p.pptx")
|
||||
|
||||
elements = partition_pptx(filename=filename)
|
||||
chunk_elements = partition_pptx(filename, chunking_strategy="by_title")
|
||||
chunks = chunk_by_title(elements)
|
||||
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_partition_pptx_title_shape_detection(tmp_path: pathlib.Path):
|
||||
"""This tests if the title attribute of a shape is correctly categorized as a title"""
|
||||
filename = str(tmp_path / "test-title-shape.pptx")
|
||||
|
||||
# create a fake PowerPoint presentation with a slide containing a title shape
|
||||
prs = pptx.Presentation()
|
||||
slide = prs.slides.add_slide(prs.slide_layouts[0])
|
||||
title_shape = slide.shapes.title
|
||||
assert title_shape is not None
|
||||
title_shape.text = (
|
||||
"This is a title, it's a bit long so we can make sure it's not narrative text"
|
||||
)
|
||||
title_shape.text_frame.add_paragraph().text = "this is a subtitle"
|
||||
|
||||
prs.save(filename)
|
||||
|
||||
# partition the PowerPoint presentation and get the first element
|
||||
elements = partition_pptx(filename)
|
||||
title = elements[0]
|
||||
subtitle = elements[1]
|
||||
|
||||
# assert that the first line is a title and has the correct text and depth
|
||||
assert isinstance(title, Title)
|
||||
assert (
|
||||
title.text == "This is a title, it's a bit long so we can make sure it's not narrative text"
|
||||
)
|
||||
assert title.metadata.category_depth == 0
|
||||
|
||||
# assert that the first line is the subtitle and has the correct text and depth
|
||||
assert isinstance(subtitle, Title)
|
||||
assert subtitle.text == "this is a subtitle"
|
||||
assert subtitle.metadata.category_depth == 1
|
||||
|
||||
|
||||
def test_partition_pptx_level_detection(tmp_path: pathlib.Path):
|
||||
"""This tests if the level attribute of a paragraph is correctly set as the category depth"""
|
||||
filename = str(tmp_path / "test-category-depth.pptx")
|
||||
|
||||
prs = pptx.Presentation()
|
||||
blank_slide_layout = prs.slide_layouts[1]
|
||||
|
||||
slide = prs.slides.add_slide(blank_slide_layout)
|
||||
shapes = slide.shapes
|
||||
title_shape = shapes.title
|
||||
assert title_shape is not None
|
||||
title_shape.text = (
|
||||
"This is a title, it's a bit long so we can make sure it's not narrative text"
|
||||
)
|
||||
body_shape = shapes.placeholders[1]
|
||||
|
||||
tf = body_shape.text_frame
|
||||
tf.text = "this is the root level bullet"
|
||||
|
||||
p = tf.add_paragraph()
|
||||
p.text = "this is the level 1 bullet"
|
||||
p.level = 1
|
||||
|
||||
p = tf.add_paragraph()
|
||||
p.text = "this is the level 2 bullet"
|
||||
p.level = 2
|
||||
|
||||
prs.slides[0].shapes
|
||||
|
||||
prs.save(filename)
|
||||
|
||||
# partition the PowerPoint presentation and get the first element
|
||||
elements = partition_pptx(filename)
|
||||
|
||||
# NOTE(newelh) - python_pptx does not create full bullet xml, so unstructured will
|
||||
# not detect the paragraphs as bullets. This is fine for now, as
|
||||
# the level attribute is still set correctly, and what we're testing here
|
||||
test_cases = [
|
||||
(0, Title, "This is a title, it's a bit long so we can make sure it's not narrative text"),
|
||||
(0, NarrativeText, "this is the root level bullet"),
|
||||
(1, NarrativeText, "this is the level 1 bullet"),
|
||||
(2, NarrativeText, "this is the level 2 bullet"),
|
||||
]
|
||||
|
||||
for element, test_case in zip(elements, test_cases):
|
||||
assert element.text == test_case[2], f"expected {test_case[2]}, got {element.text}"
|
||||
assert isinstance(
|
||||
element,
|
||||
test_case[1],
|
||||
), f"expected {test_case[1]}, got {type(element).__name__} for {element.text}"
|
||||
assert (
|
||||
element.metadata.category_depth == test_case[0]
|
||||
), f"expected {test_case[0]}, got {element.metadata.category_depth} for {element.text}"
|
||||
|
||||
|
||||
def test_partition_pptx_hierarchy_sample_document():
|
||||
"""This tests if the hierarchy of the sample document is correctly detected"""
|
||||
elements = partition_pptx(example_doc_path("sample-presentation.pptx"))
|
||||
|
||||
test_cases = [
|
||||
(0, None, "b2859226ba1f9243fb3f1b2ace889f43"),
|
||||
(1, "b2859226ba1f9243fb3f1b2ace889f43", "d13f8827e94541c8b818b0df8f942526"),
|
||||
(None, None, "cbb95b030de22979af6bfa42969c8202"),
|
||||
(0, None, "e535f799d1f0e79d6777efa873a16ce1"),
|
||||
(0, "e535f799d1f0e79d6777efa873a16ce1", "f02bbfb417ad60daa2ba35080e96262f"),
|
||||
(0, "e535f799d1f0e79d6777efa873a16ce1", "414dfce72ea53cd4649176af0d62a4c1"),
|
||||
(1, "414dfce72ea53cd4649176af0d62a4c1", "3d45a95c79473a07db4edca5534a7c49"),
|
||||
(1, "414dfce72ea53cd4649176af0d62a4c1", "a33333f527851f700ca175acd04b8a2c"),
|
||||
(2, "a33333f527851f700ca175acd04b8a2c", "6f1b87689e4da2b0fb865bc5f92d5702"),
|
||||
(0, "e535f799d1f0e79d6777efa873a16ce1", "3f58e0be3b8e8b15cba7adc4eae68586"),
|
||||
(None, None, "e5de1b503e64da424fb7d8113371e16d"),
|
||||
(0, None, "8319096532fe2e55f66c491ea8313150"),
|
||||
(0, "8319096532fe2e55f66c491ea8313150", "17a7e78277ab131a627cb4538bab7390"),
|
||||
(0, "8319096532fe2e55f66c491ea8313150", "41a9e1d0390f4edd77181142ceae51bc"),
|
||||
(1, "41a9e1d0390f4edd77181142ceae51bc", "cbbc78ef38a035fd66f7b030dcf12f66"),
|
||||
(1, "41a9e1d0390f4edd77181142ceae51bc", "2a551e3cbe67561debe0da262a294f24"),
|
||||
(2, "2a551e3cbe67561debe0da262a294f24", "7a121a056eedb11ac8804d6fd17afc0c"),
|
||||
(0, "8319096532fe2e55f66c491ea8313150", "a24a3caf9853702cb73daae23020b7b4"),
|
||||
(0, "8319096532fe2e55f66c491ea8313150", "18367f334b5c8c4602ea413ab68ac35b"),
|
||||
(0, "8319096532fe2e55f66c491ea8313150", "7f647b1f0f20c3db40c36ab57d9a5550"),
|
||||
(1, "7f647b1f0f20c3db40c36ab57d9a5550", "591c24b41b53aba873188a0881d10961"),
|
||||
(1, "7f647b1f0f20c3db40c36ab57d9a5550", "6ec455f5f19782facf184886876c9a66"),
|
||||
(2, "6ec455f5f19782facf184886876c9a66", "5614b00c3f6bff23ebba1360e10f6428"),
|
||||
(0, "8319096532fe2e55f66c491ea8313150", "2f57a8d4182e6fd5bd5842b0a2d9841b"),
|
||||
(None, None, "4120066d251ba675ade42e8a167ca61f"),
|
||||
(None, None, "efb9d74b4f8be6308c9a9006da994e12"),
|
||||
(0, None, "fd08cacbaddafee5cbacc02528536ee5"),
|
||||
]
|
||||
|
||||
# Zip the test cases with the elements
|
||||
for element, test_case in zip(elements, test_cases):
|
||||
expected_depth, expected_parent_id, expected_id = test_case
|
||||
assert element.metadata.category_depth == expected_depth
|
||||
assert element.metadata.parent_id == expected_parent_id
|
||||
assert element.id == expected_id
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# MODULE-LEVEL FIXTURES
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def opts_args() -> dict[str, Any]:
|
||||
"""All default arguments for `_XlsxPartitionerOptions`.
|
||||
|
||||
Individual argument values can be changed to suit each test. Makes construction of opts more
|
||||
compact for testing purposes.
|
||||
"""
|
||||
return {
|
||||
"file": None,
|
||||
"file_path": None,
|
||||
"include_page_breaks": True,
|
||||
"include_slide_notes": False,
|
||||
"infer_table_structure": True,
|
||||
"strategy": "fast",
|
||||
}
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# ISOLATED UNIT TESTS
|
||||
# ================================================================================================
|
||||
# These test components used by `partition_pptx()` in isolation such that all edge cases can be
|
||||
# exercised.
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class DescribePptxPartitionerOptions:
|
||||
"""Unit-test suite for `unstructured.partition.xlsx.PptxPartitionerOptions` objects."""
|
||||
|
||||
@pytest.mark.parametrize("arg_value", [True, False])
|
||||
def it_knows_whether_to_emit_PageBreak_elements_as_part_of_the_output_element_stream(
|
||||
self, arg_value: bool, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["include_page_breaks"] = arg_value
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.include_page_breaks is arg_value
|
||||
|
||||
@pytest.mark.parametrize("arg_value", [True, False])
|
||||
def it_knows_whether_to_partition_content_found_in_slide_notes(
|
||||
self, arg_value: bool, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["include_slide_notes"] = arg_value
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.include_slide_notes is arg_value
|
||||
|
||||
@pytest.mark.parametrize("arg_value", [True, False])
|
||||
def it_knows_whether_to_include_text_as_html_in_Table_metadata(
|
||||
self, arg_value: bool, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["infer_table_structure"] = arg_value
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.infer_table_structure is arg_value
|
||||
|
||||
# -- .increment_page_number() ----------------
|
||||
|
||||
def it_generates_a_PageBreak_element_when_the_page_number_is_incremented(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
# -- move to the first slide --
|
||||
list(opts.increment_page_number())
|
||||
|
||||
page_break_iter = opts.increment_page_number()
|
||||
|
||||
assert isinstance(next(page_break_iter, None), PageBreak)
|
||||
assert opts.page_number == 2
|
||||
with pytest.raises(StopIteration):
|
||||
next(page_break_iter)
|
||||
|
||||
def but_it_does_not_generate_a_PageBreak_element_for_the_first_slide(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
page_break_iter = opts.increment_page_number()
|
||||
|
||||
with pytest.raises(StopIteration):
|
||||
next(page_break_iter)
|
||||
assert opts.page_number == 1
|
||||
|
||||
def and_it_does_not_generate_a_PageBreak_element_when_include_page_breaks_option_is_off(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["include_page_breaks"] = False
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
# -- move to the first slide --
|
||||
list(opts.increment_page_number())
|
||||
|
||||
page_break_iter = opts.increment_page_number()
|
||||
|
||||
with pytest.raises(StopIteration):
|
||||
next(page_break_iter)
|
||||
assert opts.page_number == 2
|
||||
|
||||
# -- .last_modified --------------------------
|
||||
|
||||
def it_gets_last_modified_from_the_filesystem_when_a_path_is_provided(
|
||||
self, opts_args: dict[str, Any], get_last_modified_date_: Mock
|
||||
):
|
||||
opts_args["file_path"] = "a/b/spreadsheet.pptx"
|
||||
get_last_modified_date_.return_value = "2024-04-02T20:32:35"
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
last_modified = opts.last_modified
|
||||
|
||||
get_last_modified_date_.assert_called_once_with("a/b/spreadsheet.pptx")
|
||||
assert last_modified == "2024-04-02T20:32:35"
|
||||
|
||||
def and_it_falls_back_to_None_for_the_last_modified_date_when_no_path_is_provided(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
file = io.BytesIO(b"abcdefg")
|
||||
opts_args["file"] = file
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
last_modified = opts.last_modified
|
||||
|
||||
assert last_modified is None
|
||||
|
||||
# -- .metadata_file_path ---------------------
|
||||
|
||||
@pytest.mark.parametrize("file_path", ["u/v/w.pptx", None])
|
||||
def it_uses_the_filename_argument_when_provided(
|
||||
self, file_path: str | None, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["file_path"] = file_path
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.metadata_file_path == file_path
|
||||
|
||||
# -- .page_number ----------------------------
|
||||
|
||||
def it_keeps_track_of_the_page_number(self, opts_args: dict[str, Any]):
|
||||
"""In PPTX, page-number is the slide number."""
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.page_number == 0
|
||||
list(opts.increment_page_number())
|
||||
assert opts.page_number == 1
|
||||
list(opts.increment_page_number())
|
||||
assert opts.page_number == 2
|
||||
|
||||
def it_assigns_the_correct_page_number_when_starting_page_number_is_given(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
opts = PptxPartitionerOptions(**opts_args, starting_page_number=3)
|
||||
# -- move to the "first" slide --
|
||||
list(opts.increment_page_number())
|
||||
|
||||
table_metadata = opts.table_metadata(text_as_html="<table><tr/></table>")
|
||||
text_metadata = opts.text_metadata()
|
||||
|
||||
assert isinstance(table_metadata, ElementMetadata)
|
||||
assert isinstance(text_metadata, ElementMetadata)
|
||||
assert text_metadata.page_number == 3
|
||||
assert table_metadata.page_number == 3
|
||||
|
||||
# -- .pptx_file ------------------------------
|
||||
|
||||
def it_uses_the_path_to_open_the_presentation_when_file_path_is_provided(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["file_path"] = "l/m/n.pptx"
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.pptx_file == "l/m/n.pptx"
|
||||
|
||||
def and_it_uses_a_BytesIO_file_to_replaces_a_SpooledTemporaryFile_provided(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
with tempfile.SpooledTemporaryFile() as spooled_temp_file:
|
||||
spooled_temp_file.write(b"abcdefg")
|
||||
opts_args["file"] = spooled_temp_file
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
pptx_file = opts.pptx_file
|
||||
|
||||
assert pptx_file is not spooled_temp_file
|
||||
assert isinstance(pptx_file, io.BytesIO)
|
||||
assert pptx_file.getvalue() == b"abcdefg"
|
||||
|
||||
def and_it_uses_the_provided_file_directly_when_not_a_SpooledTemporaryFile(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
file = io.BytesIO(b"abcdefg")
|
||||
opts_args["file"] = file
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
pptx_file = opts.pptx_file
|
||||
|
||||
assert pptx_file is file
|
||||
assert isinstance(pptx_file, io.BytesIO)
|
||||
assert pptx_file.getvalue() == b"abcdefg"
|
||||
|
||||
def but_it_raises_ValueError_when_neither_a_file_path_or_file_is_provided(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
with pytest.raises(ValueError, match="No PPTX document specified, either `filename` or "):
|
||||
opts.pptx_file
|
||||
|
||||
# -- .strategy -------------------------------
|
||||
|
||||
@pytest.mark.parametrize("arg_value", ["fast", "hi_res"])
|
||||
def it_knows_which_partitioning_strategy_to_use(
|
||||
self, arg_value: str, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["strategy"] = arg_value
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.strategy == arg_value
|
||||
|
||||
# -- .table_metadata -------------------------
|
||||
|
||||
def it_can_create_table_metadata(
|
||||
self, last_modified_prop_: Mock, metadata_file_path_prop_: Mock, opts_args: dict[str, Any]
|
||||
):
|
||||
metadata_file_path_prop_.return_value = "d/e/f.pptx"
|
||||
last_modified_prop_.return_value = "2024-04-02T19:51:55"
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
# -- move to the first slide --
|
||||
list(opts.increment_page_number())
|
||||
|
||||
metadata = opts.table_metadata(text_as_html="<table><tr/></table>")
|
||||
|
||||
assert isinstance(metadata, ElementMetadata)
|
||||
assert metadata.filename == "f.pptx"
|
||||
assert metadata.last_modified == "2024-04-02T19:51:55"
|
||||
assert metadata.page_number == 1
|
||||
assert metadata.text_as_html == "<table><tr/></table>"
|
||||
|
||||
# -- .text_metadata -------------------------
|
||||
|
||||
def it_can_create_text_metadata(
|
||||
self, last_modified_prop_: Mock, metadata_file_path_prop_: Mock, opts_args: dict[str, Any]
|
||||
):
|
||||
metadata_file_path_prop_.return_value = "d/e/f.pptx"
|
||||
last_modified_prop_.return_value = "2024-04-02T19:56:40"
|
||||
opts = PptxPartitionerOptions(**opts_args)
|
||||
# -- move to the first slide --
|
||||
list(opts.increment_page_number())
|
||||
|
||||
metadata = opts.text_metadata(category_depth=2)
|
||||
|
||||
assert isinstance(metadata, ElementMetadata)
|
||||
assert metadata.filename == "f.pptx"
|
||||
assert metadata.last_modified == "2024-04-02T19:56:40"
|
||||
assert metadata.page_number == 1
|
||||
assert metadata.category_depth == 2
|
||||
|
||||
# -- fixtures --------------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture()
|
||||
def get_last_modified_date_(self, request: FixtureRequest):
|
||||
return function_mock(request, "unstructured.partition.pptx.get_last_modified_date")
|
||||
|
||||
@pytest.fixture()
|
||||
def last_modified_prop_(self, request: FixtureRequest):
|
||||
return property_mock(request, PptxPartitionerOptions, "last_modified")
|
||||
|
||||
@pytest.fixture()
|
||||
def metadata_file_path_prop_(self, request: FixtureRequest):
|
||||
return property_mock(request, PptxPartitionerOptions, "metadata_file_path")
|
||||
@@ -0,0 +1,140 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.unit_utils import (
|
||||
assert_round_trips_through_JSON,
|
||||
example_doc_path,
|
||||
find_text_in_elements,
|
||||
)
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import Title
|
||||
from unstructured.partition.rst import partition_rst
|
||||
|
||||
|
||||
def test_partition_rst_from_filename():
|
||||
elements = partition_rst(example_doc_path("README.rst"))
|
||||
assert elements[0] == Title("Example Docs")
|
||||
|
||||
|
||||
def test_partition_rst_from_file():
|
||||
with open(example_doc_path("README.rst"), "rb") as f:
|
||||
elements = partition_rst(file=f)
|
||||
|
||||
assert elements[0] == Title("Example Docs")
|
||||
|
||||
|
||||
# -- .metadata.filename --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_rst_from_filename_gets_filename_from_filename_arg():
|
||||
elements = partition_rst(example_doc_path("README.rst"))
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename == "README.rst" for e in elements)
|
||||
|
||||
|
||||
def test_partition_rst_from_file_gets_filename_None():
|
||||
with open(example_doc_path("README.rst"), "rb") as f:
|
||||
elements = partition_rst(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_rst_from_filename_prefers_metadata_filename():
|
||||
elements = partition_rst(example_doc_path("README.rst"), metadata_filename="orig-name.rst")
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(element.metadata.filename == "orig-name.rst" for element in elements)
|
||||
|
||||
|
||||
def test_partition_rst_from_file_prefers_metadata_filename():
|
||||
with open(example_doc_path("README.rst"), "rb") as f:
|
||||
elements = partition_rst(file=f, metadata_filename="orig-name.rst")
|
||||
|
||||
assert all(e.metadata.filename == "orig-name.rst" for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.filetype --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_rst_gets_the_RST_MIME_type_in_metadata_filetype():
|
||||
RST_MIME_TYPE = "text/x-rst"
|
||||
elements = partition_rst(example_doc_path("README.rst"))
|
||||
assert all(e.metadata.filetype == RST_MIME_TYPE for e in elements), (
|
||||
f"Expected all elements to have '{RST_MIME_TYPE}' as their filetype, but got:"
|
||||
f" {repr(elements[0].metadata.filetype)}"
|
||||
)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_rst_from_file_path_gets_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2024-06-14T16:01:29"
|
||||
mocker.patch(
|
||||
"unstructured.partition.rst.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_rst(example_doc_path("README.rst"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_rst_from_file_path_prefers_metadata_last_modified(mocker: MockFixture):
|
||||
filesystem_last_modified = "2024-06-14T16:01:29"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.rst.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_rst(
|
||||
example_doc_path("README.rst"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_rst_with_json():
|
||||
elements = partition_rst(example_doc_path("README.rst"))
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
def test_add_chunking_strategy_on_partition_rst():
|
||||
file_path = example_doc_path("README.rst")
|
||||
elements = partition_rst(file_path)
|
||||
chunk_elements = partition_rst(file_path, chunking_strategy="by_title")
|
||||
chunks = chunk_by_title(elements)
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_partition_rst_element_metadata_has_languages():
|
||||
elements = partition_rst(example_doc_path("README.rst"))
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_rst_respects_detect_language_per_element():
|
||||
elements = partition_rst(
|
||||
example_doc_path("language-docs/eng_spa_mult.rst"), detect_language_per_element=True
|
||||
)
|
||||
langs = [element.metadata.languages for element in elements]
|
||||
assert langs == [["eng"], ["spa", "eng"], ["eng"], ["eng"], ["spa"]]
|
||||
|
||||
|
||||
def test_rst_wont_include_external_files():
|
||||
# Make sure our import file is in place (otherwise the import fails silently and test passes)
|
||||
assert Path(example_doc_path("file_we_dont_want_imported")).exists()
|
||||
elements = partition_rst(example_doc_path("README-w-include.rst"))
|
||||
# The partition should contain some elements
|
||||
assert elements
|
||||
# We find something we expect to find from file we partitioned directly
|
||||
assert find_text_in_elements("instructions", elements)
|
||||
# But we don't find something from the file included within the file we partitioned directly
|
||||
assert not find_text_in_elements("wombat", elements)
|
||||
@@ -0,0 +1,129 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.unit_utils import assert_round_trips_through_JSON, example_doc_path
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import Table, Title
|
||||
from unstructured.partition.rtf import partition_rtf
|
||||
|
||||
|
||||
def test_partition_rtf_from_filename():
|
||||
elements = partition_rtf(example_doc_path("fake-doc.rtf"))
|
||||
|
||||
assert len(elements) > 0
|
||||
assert elements[0] == Title("My First Heading")
|
||||
assert elements[-1] == Table(
|
||||
text="Column 1 Column 2 Row 1, Cell 1 Row 1, Cell 2 Row 2, Cell 1 Row 2, Cell 2"
|
||||
)
|
||||
|
||||
|
||||
def test_partition_rtf_from_file():
|
||||
with open(example_doc_path("fake-doc.rtf"), "rb") as f:
|
||||
elements = partition_rtf(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert elements[0] == Title("My First Heading")
|
||||
|
||||
|
||||
# -- .metadata.filename --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_rtf_from_filename_gets_filename_from_filename_arg():
|
||||
elements = partition_rtf(example_doc_path("fake-doc.rtf"))
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename == "fake-doc.rtf" for e in elements)
|
||||
|
||||
|
||||
def test_partition_rtf_from_file_gets_filename_None():
|
||||
with open(example_doc_path("fake-doc.rtf"), "rb") as f:
|
||||
elements = partition_rtf(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_rtf_from_filename_prefers_metadata_filename():
|
||||
elements = partition_rtf(example_doc_path("fake-doc.rtf"), metadata_filename="orig-name.rtf")
|
||||
|
||||
assert len(elements) > 0
|
||||
assert all(element.metadata.filename == "orig-name.rtf" for element in elements)
|
||||
|
||||
|
||||
def test_partition_rtf_from_file_prefers_metadata_filename():
|
||||
with open(example_doc_path("fake-doc.rtf"), "rb") as f:
|
||||
elements = partition_rtf(file=f, metadata_filename="orig-name.rtf")
|
||||
|
||||
assert all(e.metadata.filename == "orig-name.rtf" for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.filetype --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_rtf_gets_the_RTF_MIME_type_in_metadata_filetype():
|
||||
RTF_MIME_TYPE = "text/rtf"
|
||||
elements = partition_rtf(example_doc_path("fake-doc.rtf"))
|
||||
assert all(e.metadata.filetype == RTF_MIME_TYPE for e in elements), (
|
||||
f"Expected all elements to have '{RTF_MIME_TYPE}' as their filetype, but got:"
|
||||
f" {repr(elements[0].metadata.filetype)}"
|
||||
)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_rtf_pulls_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2024-06-14T16:01:29"
|
||||
mocker.patch(
|
||||
"unstructured.partition.rtf.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_rtf("example-docs/fake-doc.rtf")
|
||||
|
||||
assert elements[0].metadata.last_modified == filesystem_last_modified
|
||||
|
||||
|
||||
def test_partition_rtf_prefers_metadata_last_modified(mocker: MockFixture):
|
||||
metadata_last_modified = "2024-06-14T16:01:29"
|
||||
mocker.patch(
|
||||
"unstructured.partition.rtf.get_last_modified_date", return_value="2029-07-05T09:24:28"
|
||||
)
|
||||
|
||||
elements = partition_rtf(
|
||||
"example-docs/fake-doc.rtf", metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# -- other ---------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_rtf_with_json():
|
||||
elements = partition_rtf(filename=example_doc_path("fake-doc.rtf"))
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
def test_add_chunking_strategy_on_partition_rtf():
|
||||
file_path = example_doc_path("fake-doc.rtf")
|
||||
elements = partition_rtf(filename=file_path)
|
||||
|
||||
chunk_elements = partition_rtf(file_path, chunking_strategy="by_title")
|
||||
chunks = chunk_by_title(elements)
|
||||
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_partition_rtf_element_metadata_has_languages():
|
||||
filename = "example-docs/fake-doc.rtf"
|
||||
elements = partition_rtf(filename=filename)
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_rtf_respects_detect_language_per_element():
|
||||
filename = "example-docs/language-docs/eng_spa_mult.rtf"
|
||||
elements = partition_rtf(filename=filename, detect_language_per_element=True)
|
||||
langs = [element.metadata.languages for element in elements]
|
||||
assert langs == [["eng"], ["spa", "eng"], ["eng"], ["eng"], ["spa"]]
|
||||
@@ -0,0 +1,130 @@
|
||||
import pytest
|
||||
|
||||
from test_unstructured.unit_utils import example_doc_path
|
||||
from unstructured.documents.elements import Text
|
||||
from unstructured.partition import pdf, strategies
|
||||
from unstructured.partition.utils.constants import PartitionStrategy
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"strategy",
|
||||
[
|
||||
PartitionStrategy.AUTO,
|
||||
PartitionStrategy.FAST,
|
||||
PartitionStrategy.OCR_ONLY,
|
||||
PartitionStrategy.HI_RES,
|
||||
],
|
||||
)
|
||||
def test_validate_strategy(strategy):
|
||||
# Nothing should raise for a valid strategy
|
||||
strategies.validate_strategy(strategy=strategy)
|
||||
|
||||
|
||||
def test_validate_strategy_raises_for_fast_strategy():
|
||||
with pytest.raises(ValueError):
|
||||
strategies.validate_strategy(strategy=PartitionStrategy.FAST, is_image=True)
|
||||
|
||||
|
||||
def test_validate_strategy_raises_for_bad_strategy():
|
||||
with pytest.raises(ValueError):
|
||||
strategies.validate_strategy("totally_guess_the_text")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "from_file", "expected"),
|
||||
[
|
||||
("layout-parser-paper-fast.pdf", True, True),
|
||||
("copy-protected.pdf", True, True),
|
||||
("loremipsum-flat.pdf", True, False),
|
||||
("layout-parser-paper-fast.pdf", False, True),
|
||||
("copy-protected.pdf", False, True),
|
||||
("loremipsum-flat.pdf", False, False),
|
||||
],
|
||||
)
|
||||
def test_is_pdf_text_extractable(filename, from_file, expected):
|
||||
filename = example_doc_path(f"pdf/{filename}")
|
||||
|
||||
if from_file:
|
||||
with open(filename, "rb") as f:
|
||||
extracted_elements = pdf.extractable_elements(file=f)
|
||||
else:
|
||||
extracted_elements = pdf.extractable_elements(filename=filename)
|
||||
|
||||
pdf_text_extractable = any(
|
||||
isinstance(el, Text) and el.text.strip()
|
||||
for page_elements in extracted_elements
|
||||
for el in page_elements
|
||||
)
|
||||
|
||||
assert pdf_text_extractable is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("pdf_text_extractable", "infer_table_structure"),
|
||||
[
|
||||
(True, True),
|
||||
(False, True),
|
||||
(True, False),
|
||||
(False, False),
|
||||
],
|
||||
)
|
||||
def test_determine_pdf_or_image_fast_strategy(pdf_text_extractable, infer_table_structure):
|
||||
strategy = strategies.determine_pdf_or_image_strategy(
|
||||
strategy=PartitionStrategy.FAST,
|
||||
pdf_text_extractable=pdf_text_extractable,
|
||||
infer_table_structure=infer_table_structure,
|
||||
)
|
||||
assert strategy == PartitionStrategy.FAST
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
(
|
||||
"pdf_text_extractable",
|
||||
"infer_table_structure",
|
||||
"extract_images_in_pdf",
|
||||
"extract_image_block_types",
|
||||
"expected",
|
||||
),
|
||||
[
|
||||
(True, True, True, ["Image"], PartitionStrategy.HI_RES),
|
||||
(True, True, True, [], PartitionStrategy.HI_RES),
|
||||
(True, True, False, ["Image"], PartitionStrategy.HI_RES),
|
||||
(True, True, False, [], PartitionStrategy.HI_RES),
|
||||
(True, False, True, ["Image"], PartitionStrategy.HI_RES),
|
||||
(True, False, True, [], PartitionStrategy.HI_RES),
|
||||
(True, False, False, ["Image"], PartitionStrategy.HI_RES),
|
||||
(True, False, False, [], PartitionStrategy.FAST),
|
||||
(False, True, True, ["Image"], PartitionStrategy.HI_RES),
|
||||
(False, True, True, [], PartitionStrategy.HI_RES),
|
||||
(False, True, False, ["Image"], PartitionStrategy.HI_RES),
|
||||
(False, True, False, [], PartitionStrategy.HI_RES),
|
||||
(False, False, True, ["Image"], PartitionStrategy.HI_RES),
|
||||
(False, False, True, [], PartitionStrategy.HI_RES),
|
||||
(False, False, False, ["Image"], PartitionStrategy.HI_RES),
|
||||
(False, False, False, [], PartitionStrategy.OCR_ONLY),
|
||||
],
|
||||
)
|
||||
def test_determine_pdf_auto_strategy(
|
||||
pdf_text_extractable,
|
||||
infer_table_structure,
|
||||
extract_images_in_pdf,
|
||||
extract_image_block_types,
|
||||
expected,
|
||||
):
|
||||
strategy = strategies.determine_pdf_or_image_strategy(
|
||||
strategy=PartitionStrategy.AUTO,
|
||||
is_image=False,
|
||||
pdf_text_extractable=pdf_text_extractable,
|
||||
infer_table_structure=infer_table_structure,
|
||||
extract_images_in_pdf=extract_images_in_pdf,
|
||||
extract_image_block_types=extract_image_block_types,
|
||||
)
|
||||
assert strategy == expected
|
||||
|
||||
|
||||
def test_determine_image_auto_strategy():
|
||||
strategy = strategies.determine_pdf_or_image_strategy(
|
||||
strategy=PartitionStrategy.AUTO,
|
||||
is_image=True,
|
||||
)
|
||||
assert strategy == PartitionStrategy.HI_RES
|
||||
@@ -0,0 +1,446 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from typing import Optional, Type
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from test_unstructured.unit_utils import assert_round_trips_through_JSON, example_doc_path
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.cleaners.core import group_broken_paragraphs
|
||||
from unstructured.documents.elements import Address, ListItem, NarrativeText, Title
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.partition.text import partition_text
|
||||
from unstructured.partition.utils.constants import UNSTRUCTURED_INCLUDE_DEBUG_METADATA
|
||||
|
||||
EXPECTED_OUTPUT = [
|
||||
NarrativeText(text="This is a test document to use for unit tests."),
|
||||
Address(text="Doylestown, PA 18901"),
|
||||
Title(text="Important points:"),
|
||||
ListItem(text="Hamburgers are delicious"),
|
||||
ListItem(text="Dogs are the best"),
|
||||
ListItem(text="I love fuzzy blankets"),
|
||||
]
|
||||
|
||||
MIN_MAX_TEXT = """This is a story. This is a story that doesn't matter
|
||||
because it is just being used as an example. Hi. Hello. Howdy. Hola.
|
||||
The example is simple and repetitive and long and somewhat boring,
|
||||
but it serves a purpose. End.""".replace(
|
||||
"\n",
|
||||
"",
|
||||
)
|
||||
|
||||
SHORT_PARAGRAPHS = """This is a story.
|
||||
|
||||
This is a story that doesn't matter because it is just being used as an example.
|
||||
|
||||
Hi.
|
||||
|
||||
Hello.
|
||||
|
||||
Howdy.
|
||||
|
||||
Hola.
|
||||
|
||||
The example is simple and repetitive and long and somewhat boring, but it serves a purpose.
|
||||
|
||||
End.
|
||||
"""
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "encoding"),
|
||||
[
|
||||
("fake-text.txt", "utf-8"),
|
||||
("fake-text.txt", None),
|
||||
("fake-text-utf-16-be.txt", "utf-16-be"),
|
||||
],
|
||||
)
|
||||
def test_partition_text_from_filename(filename: str, encoding: Optional[str]):
|
||||
elements = partition_text(example_doc_path(filename), encoding=encoding)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert elements == EXPECTED_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename == filename
|
||||
if UNSTRUCTURED_INCLUDE_DEBUG_METADATA:
|
||||
assert {element.metadata.detection_origin for element in elements} == {"text"}
|
||||
|
||||
|
||||
def test_partition_text_from_filename_with_metadata_filename():
|
||||
elements = partition_text(
|
||||
example_doc_path("fake-text.txt"), encoding="utf-8", metadata_filename="test"
|
||||
)
|
||||
|
||||
assert elements == EXPECTED_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "test"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename",
|
||||
["fake-text-utf-16.txt", "fake-text-utf-16-le.txt", "fake-text-utf-32.txt"],
|
||||
)
|
||||
def test_partition_text_from_filename_default_encoding(filename: str):
|
||||
elements = partition_text(example_doc_path(filename))
|
||||
|
||||
assert len(elements) > 0
|
||||
assert elements == EXPECTED_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename == filename
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "encoding", "error"),
|
||||
[
|
||||
("fake-text.txt", "utf-16", UnicodeDecodeError),
|
||||
("fake-text-utf-16-be.txt", "utf-16", UnicodeError),
|
||||
],
|
||||
)
|
||||
def test_partition_text_from_filename_raises_econding_error(
|
||||
filename: str,
|
||||
encoding: Optional[str],
|
||||
error: Type[BaseException],
|
||||
):
|
||||
with pytest.raises(error):
|
||||
filename = example_doc_path(filename)
|
||||
partition_text(filename=filename, encoding=encoding)
|
||||
|
||||
|
||||
def test_partition_text_from_file():
|
||||
with open(example_doc_path("fake-text.txt"), "rb") as f:
|
||||
elements = partition_text(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert elements == EXPECTED_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename is None
|
||||
|
||||
|
||||
def test_partition_text_from_file_with_metadata_filename():
|
||||
filename = example_doc_path("fake-text.txt")
|
||||
with open(filename, "rb") as f:
|
||||
elements = partition_text(file=f, metadata_filename="test")
|
||||
assert len(elements) > 0
|
||||
assert elements == EXPECTED_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename == "test"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename",
|
||||
["fake-text-utf-16.txt", "fake-text-utf-16-le.txt", "fake-text-utf-32.txt"],
|
||||
)
|
||||
def test_partition_text_from_file_default_encoding(filename: str):
|
||||
with open(example_doc_path(filename), "rb") as f:
|
||||
elements = partition_text(file=f)
|
||||
assert len(elements) > 0
|
||||
assert elements == EXPECTED_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename is None
|
||||
|
||||
|
||||
def test_partition_text_from_bytes_file():
|
||||
with open(example_doc_path("fake-text.txt"), "rb") as f:
|
||||
elements = partition_text(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert elements == EXPECTED_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename",
|
||||
["fake-text-utf-16.txt", "fake-text-utf-16-le.txt", "fake-text-utf-32.txt"],
|
||||
)
|
||||
def test_partition_text_from_bytes_file_default_encoding(filename: str):
|
||||
with open(example_doc_path(filename), "rb") as f:
|
||||
elements = partition_text(file=f)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert elements == EXPECTED_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename is None
|
||||
|
||||
|
||||
def test_text_partition_element_metadata_user_provided_languages():
|
||||
elements = partition_text(
|
||||
example_doc_path("book-war-and-peace-1p.txt"), strategy="fast", languages=["en"]
|
||||
)
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_text_from_text():
|
||||
with open(example_doc_path("fake-text.txt")) as f:
|
||||
text = f.read()
|
||||
|
||||
elements = partition_text(text=text)
|
||||
|
||||
assert len(elements) > 0
|
||||
assert elements == EXPECTED_OUTPUT
|
||||
for element in elements:
|
||||
assert element.metadata.filename is None
|
||||
|
||||
|
||||
def test_partition_text_from_text_works_with_empty_string():
|
||||
assert partition_text(text="") == []
|
||||
|
||||
|
||||
def test_partition_text_raises_with_none_specified():
|
||||
with pytest.raises(ValueError):
|
||||
partition_text()
|
||||
|
||||
|
||||
def test_partition_text_raises_with_too_many_specified():
|
||||
filename = example_doc_path("fake-text.txt")
|
||||
with open(filename) as f:
|
||||
text = f.read()
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
partition_text(filename=filename, text=text)
|
||||
|
||||
|
||||
def test_partition_text_captures_everything_even_with_linebreaks():
|
||||
text = """
|
||||
VERY IMPORTANT MEMO
|
||||
DOYLESTOWN, PA 18901
|
||||
"""
|
||||
elements = partition_text(text=text)
|
||||
assert elements == [
|
||||
Title(text="VERY IMPORTANT MEMO"),
|
||||
Address(text="DOYLESTOWN, PA 18901"),
|
||||
]
|
||||
for element in elements:
|
||||
assert element.metadata.filename is None
|
||||
|
||||
|
||||
def test_partition_text_groups_broken_paragraphs():
|
||||
text = (
|
||||
"The big brown fox\n"
|
||||
"was walking down the lane.\n"
|
||||
"\n"
|
||||
"At the end of the lane,\n"
|
||||
"the fox met a bear."
|
||||
)
|
||||
|
||||
elements = partition_text(text=text, paragraph_grouper=group_broken_paragraphs)
|
||||
|
||||
assert elements == [
|
||||
NarrativeText(text="The big brown fox was walking down the lane."),
|
||||
NarrativeText(text="At the end of the lane, the fox met a bear."),
|
||||
]
|
||||
for element in elements:
|
||||
assert element.metadata.filename is None
|
||||
|
||||
|
||||
def test_partition_text_splits_long_text():
|
||||
elements = partition_text(example_doc_path("norwich-city.txt"))
|
||||
assert len(elements) > 0
|
||||
assert elements[0].text.startswith("Iwan Roberts")
|
||||
assert elements[-1].text.endswith("External links")
|
||||
|
||||
|
||||
def test_partition_text_doesnt_get_page_breaks():
|
||||
text = "--------------------"
|
||||
elements = partition_text(text=text)
|
||||
assert len(elements) == 1
|
||||
assert elements[0].text == text
|
||||
assert not isinstance(elements[0], ListItem)
|
||||
|
||||
|
||||
# -- .metadata.filename --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_text_from_filename_gets_filename_metadata_from_file_path():
|
||||
elements = partition_text(example_doc_path("fake-text.txt"))
|
||||
|
||||
assert all(e.metadata.filename == "fake-text.txt" for e in elements)
|
||||
assert all(e.metadata.file_directory == example_doc_path("") for e in elements)
|
||||
|
||||
|
||||
def test_partition_text_from_file_gets_filename_metadata_None():
|
||||
with open(example_doc_path("fake-text.txt"), "rb") as f:
|
||||
elements = partition_text(file=f)
|
||||
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
assert all(e.metadata.file_directory is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_text_from_filename_prefers_metadata_filename():
|
||||
elements = partition_text(example_doc_path("fake-text.txt"), metadata_filename="a/b/c.txt")
|
||||
|
||||
assert all(e.metadata.filename == "c.txt" for e in elements)
|
||||
assert all(e.metadata.file_directory == "a/b" for e in elements)
|
||||
|
||||
|
||||
def test_partition_text_from_file_prefers_metadata_filename():
|
||||
with open(example_doc_path("fake-text.txt"), "rb") as f:
|
||||
elements = partition_text(file=f, metadata_filename="d/e/f.txt")
|
||||
|
||||
assert all(e.metadata.filename == "f.txt" for e in elements)
|
||||
assert all(e.metadata.file_directory == "d/e" for e in elements)
|
||||
|
||||
|
||||
# -- .metadata.filetype --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_text_gets_the_TXT_MIME_type_in_metadata_filetype():
|
||||
TXT_MIME_TYPE = "text/plain"
|
||||
elements = partition_text(example_doc_path("fake-text.txt"))
|
||||
assert all(e.metadata.filetype == TXT_MIME_TYPE for e in elements), (
|
||||
f"Expected all elements to have '{TXT_MIME_TYPE}' as their filetype, but got:"
|
||||
f" {repr(elements[0].metadata.filetype)}"
|
||||
)
|
||||
|
||||
|
||||
def test_partition_text_prefers_metadata_file_type():
|
||||
elements = partition_text(example_doc_path("README.md"), metadata_file_type=FileType.MD)
|
||||
assert all(e.metadata.filetype == "text/markdown" for e in elements), (
|
||||
f"Expected all elements to have 'text/markdown' as their filetype, but got:"
|
||||
f" {repr(elements[0].metadata.filetype)}"
|
||||
)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_text_from_file_path_gets_last_modified_from_filesystem(mocker: MockerFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.text.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_text(example_doc_path("fake-text.txt"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_text_from_file_gets_last_modified_None():
|
||||
with open(example_doc_path("fake-text.txt"), "rb") as f:
|
||||
elements = partition_text(file=f)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_text_from_text_gets_last_modified_None():
|
||||
with open(example_doc_path("fake-text.txt")) as f:
|
||||
text = f.read()
|
||||
|
||||
elements = partition_text(text=text)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_text_from_file_path_prefers_metadata_last_modified(mocker: MockerFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.text.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_text(
|
||||
example_doc_path("fake-text.txt"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_text_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
with open(example_doc_path("fake-text.txt"), "rb") as f:
|
||||
elements = partition_text(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_text_from_text_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
with open(example_doc_path("fake-text.txt")) as f:
|
||||
text = f.read()
|
||||
|
||||
elements = partition_text(text=text, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_Text_element_assigns_id_hashes_that_are_unique_and_deterministic():
|
||||
ids = [element.id for element in partition_text(text="hello\nhello\nhello")]
|
||||
assert ids == [
|
||||
"8657c0ec31a4cfc822f6cd4a5684cafd",
|
||||
"72aefb4a12be063ad160931fdb380163",
|
||||
"ba8c1a216ca585aecdd365a72e6124f1",
|
||||
]
|
||||
|
||||
|
||||
def test_Text_element_assings_UUID_when_unique_element_ids_is_True():
|
||||
elements = partition_text(text="hello\nhello\nhello", unique_element_ids=True)
|
||||
|
||||
for element in elements:
|
||||
assert uuid.UUID(element.id, version=4)
|
||||
|
||||
# Test that the element is JSON serializable. This should run without an error
|
||||
json.dumps(element.to_dict())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_name", "encoding"),
|
||||
[
|
||||
("fake-text.txt", "utf-8"),
|
||||
("fake-text.txt", None),
|
||||
("fake-text-utf-16-be.txt", "utf-16-be"),
|
||||
],
|
||||
)
|
||||
def test_partition_text_with_json(file_name: str, encoding: str | None):
|
||||
elements = partition_text(example_doc_path(file_name), encoding=encoding)
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
def test_add_chunking_strategy_on_partition_text():
|
||||
filename = example_doc_path("book-war-and-peace-1p.txt")
|
||||
elements = partition_text(filename=filename)
|
||||
chunk_elements = partition_text(filename, chunking_strategy="by_title")
|
||||
chunks = chunk_by_title(elements)
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_partition_text_element_metadata_has_languages():
|
||||
elements = partition_text(example_doc_path("norwich-city.txt"))
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_text_respects_detect_language_per_element():
|
||||
elements = partition_text(
|
||||
example_doc_path("language-docs/eng_spa_mult.txt"), detect_language_per_element=True
|
||||
)
|
||||
|
||||
langs = [element.metadata.languages for element in elements]
|
||||
|
||||
assert langs == [["eng"], ["spa", "eng"], ["eng"], ["eng"], ["spa"]]
|
||||
|
||||
|
||||
def test_partition_text_respects_languages_arg():
|
||||
elements = partition_text(example_doc_path("norwich-city.txt"), languages=["deu"])
|
||||
assert elements[0].metadata.languages == ["deu"]
|
||||
|
||||
|
||||
def test_partition_text_element_metadata_raises_TypeError():
|
||||
with pytest.raises(TypeError):
|
||||
partition_text(example_doc_path("norwich-city.txt"), languages="eng")
|
||||
|
||||
|
||||
def test_partition_text_detects_more_than_3_languages():
|
||||
elements = partition_text(
|
||||
example_doc_path("language-docs/UDHR_first_article_all.txt"),
|
||||
detect_language_per_element=True,
|
||||
)
|
||||
langs = [e.metadata.languages[0] for e in elements if e.metadata.languages]
|
||||
assert len(langs) > 10
|
||||
@@ -0,0 +1,324 @@
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from test_unstructured.nlp.mock_nltk import (
|
||||
mock_pos_tag,
|
||||
mock_sent_tokenize,
|
||||
mock_word_tokenize,
|
||||
)
|
||||
from unstructured.partition import text_type
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
(
|
||||
"ITEM 5(a).: MARKET FOR REGISTRANT’S COMMON EQUITY, RELATED STOCKHOLDER MATTERS AND "
|
||||
"ISSUER PURCHASES OF EQUITY SECURITIES",
|
||||
False,
|
||||
),
|
||||
(
|
||||
"Item 5(a).: Market For Registrant’s Common Equity, Related Stockholder Matters and "
|
||||
"Issuer Purchases of Equity Securities",
|
||||
False,
|
||||
),
|
||||
(
|
||||
"There is a market for registrant’s common equity, related stockholder matters and "
|
||||
"issuer purchases of equity securities.",
|
||||
True,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_headings_are_not_narrative_text(text, expected):
|
||||
assert text_type.is_possible_narrative_text(text) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("Ask the teacher for an apple.", True),
|
||||
("Ask Me About Intellectual Property", False), # Exceeds the cap threshold
|
||||
("7", False), # Fails because it is numeric
|
||||
("intellectual property", False), # Fails because it does not contain a verb
|
||||
("Dal;kdjfal adawels adfjwalsdf. Addad jaja fjawlek", False),
|
||||
("---------------Aske the teacher for an apple----------", False), # Too many non-alpha
|
||||
("", False), # Doesn't have english words # Fails because it is empty
|
||||
],
|
||||
)
|
||||
def test_is_possible_narrative_text(text, expected, monkeypatch):
|
||||
monkeypatch.setattr(text_type, "word_tokenize", mock_word_tokenize)
|
||||
monkeypatch.setattr(text_type, "pos_tag", mock_pos_tag)
|
||||
monkeypatch.setattr(text_type, "sent_tokenize", mock_sent_tokenize)
|
||||
monkeypatch.setenv("UNSTRUCTURED_LANGUAGE_CHECKS", "true")
|
||||
is_possible_narrative = text_type.is_possible_narrative_text(text, cap_threshold=0.3)
|
||||
assert is_possible_narrative is expected
|
||||
|
||||
|
||||
def test_narrative_text_language_checks():
|
||||
# NOTE(robinson) - This is true because we don't check english vocab if language checks
|
||||
# are set to False
|
||||
text = "Dal;kdjfal adawels adfjwalsdf. Addad jaja fjawlek"
|
||||
assert text_type.is_possible_narrative_text(text, language_checks=True) is False
|
||||
|
||||
|
||||
def test_text_type_handles_non_english_examples(monkeypatch):
|
||||
monkeypatch.setenv("UNSTRUCTURED_LANGUAGE_CHECKS", "true")
|
||||
narrative_text = "Я говорю по-русски. Вы тоже?"
|
||||
title = "Риски"
|
||||
|
||||
assert text_type.is_possible_narrative_text(narrative_text, languages=["eng"]) is False
|
||||
assert text_type.is_possible_narrative_text(narrative_text, languages=[]) is True
|
||||
|
||||
assert text_type.is_possible_narrative_text(title, languages=["eng"]) is False
|
||||
assert text_type.is_possible_narrative_text(title, languages=[]) is False
|
||||
|
||||
assert text_type.is_possible_title(title, languages=["eng"]) is False
|
||||
assert text_type.is_possible_title(title, languages=[]) is True
|
||||
|
||||
|
||||
def test_text_type_handles_multi_language_examples(monkeypatch):
|
||||
monkeypatch.setenv("UNSTRUCTURED_LANGUAGE_CHECKS", "true")
|
||||
narrative_text = "Я говорю по-русски. Вы тоже? 不,我不会说俄语。"
|
||||
title = "Риски (Riesgos)"
|
||||
|
||||
assert text_type.is_possible_narrative_text(narrative_text, languages=["eng"]) is False
|
||||
assert text_type.is_possible_narrative_text(narrative_text, languages=["chi", "rus"]) is True
|
||||
assert text_type.is_possible_narrative_text(narrative_text, languages=[]) is True
|
||||
|
||||
assert text_type.is_possible_narrative_text(title, languages=["eng"]) is False
|
||||
assert text_type.is_possible_narrative_text(title, languages=["spa", "rus"]) is False
|
||||
assert text_type.is_possible_narrative_text(title, languages=[]) is False
|
||||
|
||||
assert text_type.is_possible_title(title, languages=["eng"]) is False
|
||||
assert text_type.is_possible_title(title, languages=["spa", "rus"]) is True
|
||||
assert text_type.is_possible_title(title, languages=[]) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("Intellectual Property", True), # Fails because it exceeds the cap threshold
|
||||
(
|
||||
"Ask the teacher for an apple. You might a gold star.",
|
||||
False,
|
||||
), # Too many sentences
|
||||
("7", False), # Fails because it is numeric
|
||||
("", False), # Fails because it is empty
|
||||
("ITEM 1A. RISK FACTORS", True), # Two "sentences", but both are short
|
||||
("To My Dearest Friends,", False), # Ends with a comma
|
||||
("BTAR ADFJA L", False), # Doesn't have english words
|
||||
("ITEM 1A. RISK FACTORS " * 15, False), # Title is too long
|
||||
("/--------BREAK-------/", False), # Contains too many non-alpha characters
|
||||
("1.A.RISKS", True), # Tests that "RISKS" gets flagged as an english word
|
||||
("1. Unstructured Technologies", True), # Make sure we're English words :-)
|
||||
("Big/Brown/Sheet", True),
|
||||
("LOOK AT THIS IT IS CAPS BUT NOT A TITLE.", False),
|
||||
],
|
||||
)
|
||||
def test_is_possible_title(text, expected, monkeypatch):
|
||||
monkeypatch.setattr(text_type, "sent_tokenize", mock_sent_tokenize)
|
||||
monkeypatch.setattr(text_type, "word_tokenize", mock_word_tokenize)
|
||||
monkeypatch.setenv("UNSTRUCTURED_LANGUAGE_CHECKS", "true")
|
||||
assert text_type.is_possible_title(text) is expected
|
||||
|
||||
|
||||
def test_title_language_checks():
|
||||
# NOTE(robinson) - This is true because we don't check english vocab if language checks
|
||||
# are set to False
|
||||
text = "BTAR ADFJA L"
|
||||
assert text_type.is_possible_narrative_text(text, language_checks=True) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("8675309", True),
|
||||
("+1 867-5309", True),
|
||||
("2158675309", True),
|
||||
("+12158675309", True),
|
||||
("867.5309", True),
|
||||
("1-800-867-5309", True),
|
||||
("1(800)-867-5309", True),
|
||||
("Tel: 1(800)-867-5309", True),
|
||||
],
|
||||
)
|
||||
def test_contains_us_phone_number(text, expected):
|
||||
has_phone_number = text_type.contains_us_phone_number(text)
|
||||
assert has_phone_number is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("• This is a fine point!", True),
|
||||
(" • This is a fine point!", True), # Has an extra space in front of the bullet
|
||||
("‣ This is a fine point!", True),
|
||||
("⁃ This is a fine point!", True),
|
||||
("⁌ This is a fine point!", True),
|
||||
("⁍ This is a fine point!", True),
|
||||
("∙ This is a fine point!", True),
|
||||
("○ This is a fine point!", True),
|
||||
("● This is a fine point!", True),
|
||||
("◘ This is a fine point!", True),
|
||||
("◦ This is a fine point!", True),
|
||||
("☙ This is a fine point!", True),
|
||||
("❥ This is a fine point!", True),
|
||||
("❧ This is a fine point!", True),
|
||||
("⦾ This is a fine point!", True),
|
||||
("⦿ This is a fine point!", True),
|
||||
(" This is a fine point!", True),
|
||||
("* This is a fine point!", True),
|
||||
("- This is a fine point!", True),
|
||||
("This is NOT a fine point!", False), # No bullet point
|
||||
("I love morse code! ● ● ● --- ● ● ●", False), # Not at the beginning
|
||||
("----------------------------", False), # Too long
|
||||
],
|
||||
)
|
||||
def test_is_bulletized_text(text, expected):
|
||||
assert text_type.is_bulleted_text(text) is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("Ask the teacher for an apple", True),
|
||||
("Intellectual property", False),
|
||||
("THIS MESSAGE WAS APPROVED", True),
|
||||
],
|
||||
)
|
||||
def test_contains_verb(text, expected, monkeypatch):
|
||||
has_verb = text_type.contains_verb(text)
|
||||
assert has_verb is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("PARROT BEAK", True),
|
||||
("Parrot Beak", True),
|
||||
("parrot beak", True),
|
||||
("parrot!", True),
|
||||
("?parrot", True),
|
||||
("zombie?parrot", True),
|
||||
("notaWordHa 'parrot'", True),
|
||||
("notaWordHa'parrot'", False),
|
||||
('notaWordHa "parrot,"', True),
|
||||
("daljdf adlfajldj ajadfa", False),
|
||||
("BTAR ADFJA L", False),
|
||||
("Unstructured Technologies", True),
|
||||
("1.A.RISKS", True), # Test crammed together words get picked up
|
||||
("Big/Brown/Sheep", True),
|
||||
],
|
||||
)
|
||||
def test_contains_english_word(text, expected, monkeypatch):
|
||||
assert text_type.contains_english_word(text) is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("Intellectual Property in the United States", True),
|
||||
("Intellectual property helps incentivize innovation.", False),
|
||||
("THIS IS ALL CAPS. BUT IT IS TWO SENTENCES.", False),
|
||||
("LOOK AT THIS IT IS CAPS BUT NOT A TITLE.", True),
|
||||
("This Has All Caps. It's Weird But Two Sentences", False),
|
||||
("The Business Report is expected within 6 hours of closing", False),
|
||||
("", True),
|
||||
],
|
||||
)
|
||||
def test_contains_exceeds_cap_ratio(text, expected, monkeypatch):
|
||||
assert text_type.exceeds_cap_ratio(text) is expected
|
||||
|
||||
|
||||
def test_set_caps_ratio_with_environment_variable(monkeypatch):
|
||||
monkeypatch.setattr(text_type, "word_tokenize", mock_word_tokenize)
|
||||
monkeypatch.setattr(text_type, "sent_tokenize", mock_sent_tokenize)
|
||||
monkeypatch.setenv("UNSTRUCTURED_NARRATIVE_TEXT_CAP_THRESHOLD", 0.8)
|
||||
|
||||
text = "All The King's Horses. And All The King's Men."
|
||||
with patch.object(text_type, "exceeds_cap_ratio", return_value=False) as mock_exceeds:
|
||||
text_type.is_possible_narrative_text(text)
|
||||
|
||||
mock_exceeds.assert_called_once_with(text, threshold=0.8)
|
||||
|
||||
|
||||
def test_set_title_non_alpha_threshold_with_environment_variable(monkeypatch):
|
||||
monkeypatch.setattr(text_type, "word_tokenize", mock_word_tokenize)
|
||||
monkeypatch.setattr(text_type, "sent_tokenize", mock_sent_tokenize)
|
||||
monkeypatch.setenv("UNSTRUCTURED_TITLE_NON_ALPHA_THRESHOLD", 0.8)
|
||||
|
||||
text = "/--------------- All the king's horses----------------/"
|
||||
with patch.object(text_type, "under_non_alpha_ratio", return_value=False) as mock_exceeds:
|
||||
text_type.is_possible_title(text)
|
||||
|
||||
mock_exceeds.assert_called_once_with(text, threshold=0.8)
|
||||
|
||||
|
||||
def test_set_narrative_text_non_alpha_threshold_with_environment_variable(monkeypatch):
|
||||
monkeypatch.setattr(text_type, "word_tokenize", mock_word_tokenize)
|
||||
monkeypatch.setattr(text_type, "sent_tokenize", mock_sent_tokenize)
|
||||
monkeypatch.setenv("UNSTRUCTURED_NARRATIVE_TEXT_NON_ALPHA_THRESHOLD", 0.8)
|
||||
|
||||
text = "/--------------- All the king's horses----------------/"
|
||||
with patch.object(text_type, "under_non_alpha_ratio", return_value=False) as mock_exceeds:
|
||||
text_type.is_possible_narrative_text(text)
|
||||
|
||||
mock_exceeds.assert_called_once_with(text, threshold=0.8)
|
||||
|
||||
|
||||
def test_set_title_max_word_length_with_environment_variable(monkeypatch):
|
||||
monkeypatch.setattr(text_type, "word_tokenize", mock_word_tokenize)
|
||||
monkeypatch.setattr(text_type, "sent_tokenize", mock_sent_tokenize)
|
||||
monkeypatch.setenv("UNSTRUCTURED_TITLE_MAX_WORD_LENGTH", 5)
|
||||
|
||||
text = "Intellectual Property in the United States"
|
||||
assert text_type.is_possible_narrative_text(text) is False
|
||||
|
||||
|
||||
def test_sentence_count(monkeypatch):
|
||||
monkeypatch.setattr(text_type, "sent_tokenize", mock_sent_tokenize)
|
||||
text = "Hi my name is Matt. I work with Crag."
|
||||
assert text_type.sentence_count(text) == 2
|
||||
|
||||
|
||||
def test_item_titles():
|
||||
text = "ITEM 1(A). THIS IS A TITLE"
|
||||
assert text_type.sentence_count(text, 3) < 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("Doylestown, PA 18901", True),
|
||||
("DOYLESTOWN, PENNSYLVANIA, 18901", True),
|
||||
("DOYLESTOWN, PENNSYLVANIA 18901", True),
|
||||
("Doylestown, Pennsylvania 18901", True),
|
||||
(" Doylestown, Pennsylvania 18901", True),
|
||||
("The Business Report is expected within 6 hours of closing", False),
|
||||
("", False),
|
||||
],
|
||||
)
|
||||
def test_is_us_city_state_zip(text, expected):
|
||||
assert text_type.is_us_city_state_zip(text) is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("fake@gmail.com", True),
|
||||
("Fake@gmail.com", False),
|
||||
("fake.gmail.@gmail.com", True),
|
||||
("fake.gmail@.@gmail.com", False),
|
||||
(" fake@gmail.com", True),
|
||||
("fak!/$e@gmail.com", False),
|
||||
("", False),
|
||||
],
|
||||
)
|
||||
def test_is_email_address(text, expected):
|
||||
assert text_type.is_email_address(text) is expected
|
||||
|
||||
|
||||
def test_under_non_alpha_ratio_zero_divide():
|
||||
# Threw an error before changes
|
||||
text_type.under_non_alpha_ratio(" ")
|
||||
@@ -0,0 +1,161 @@
|
||||
"""Test-suite for `unstructured.partition.tsv` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockFixture
|
||||
|
||||
from test_unstructured.partition.test_constants import (
|
||||
EXPECTED_TABLE,
|
||||
EXPECTED_TABLE_WITH_EMOJI,
|
||||
EXPECTED_TEXT,
|
||||
EXPECTED_TEXT_WITH_EMOJI,
|
||||
EXPECTED_TEXT_XLSX,
|
||||
)
|
||||
from test_unstructured.unit_utils import assert_round_trips_through_JSON, example_doc_path
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import Table
|
||||
from unstructured.partition.tsv import partition_tsv
|
||||
|
||||
EXPECTED_FILETYPE = "text/tsv"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "expected_text", "expected_table"),
|
||||
[
|
||||
("stanley-cups.tsv", EXPECTED_TEXT, EXPECTED_TABLE),
|
||||
("stanley-cups-with-emoji.tsv", EXPECTED_TEXT_WITH_EMOJI, EXPECTED_TABLE_WITH_EMOJI),
|
||||
],
|
||||
)
|
||||
def test_partition_tsv_from_filename(filename: str, expected_text: str, expected_table: str):
|
||||
elements = partition_tsv(example_doc_path(filename), include_header=False)
|
||||
|
||||
table = elements[0]
|
||||
assert table.text == expected_text
|
||||
assert table.metadata.text_as_html == expected_table
|
||||
assert table.metadata.filetype == EXPECTED_FILETYPE
|
||||
assert all(e.metadata.filename == filename for e in elements)
|
||||
|
||||
|
||||
def test_partition_tsv_from_filename_with_metadata_filename():
|
||||
elements = partition_tsv(
|
||||
example_doc_path("stanley-cups.tsv"), metadata_filename="test", include_header=False
|
||||
)
|
||||
|
||||
assert elements[0].text == EXPECTED_TEXT
|
||||
assert all(e.metadata.filename == "test" for e in elements)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "expected_text", "expected_table"),
|
||||
[
|
||||
("stanley-cups.tsv", EXPECTED_TEXT, EXPECTED_TABLE),
|
||||
("stanley-cups-with-emoji.tsv", EXPECTED_TEXT_WITH_EMOJI, EXPECTED_TABLE_WITH_EMOJI),
|
||||
],
|
||||
)
|
||||
def test_partition_tsv_from_file(filename: str, expected_text: str, expected_table: str):
|
||||
with open(example_doc_path(filename), "rb") as f:
|
||||
elements = partition_tsv(file=f, include_header=False)
|
||||
|
||||
table = elements[0]
|
||||
assert isinstance(table, Table)
|
||||
assert table.text == expected_text
|
||||
assert table.metadata.text_as_html == expected_table
|
||||
assert table.metadata.filetype == EXPECTED_FILETYPE
|
||||
assert all(e.metadata.filename is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_tsv_from_file_with_metadata_filename():
|
||||
with open(example_doc_path("stanley-cups.tsv"), "rb") as f:
|
||||
elements = partition_tsv(file=f, metadata_filename="test", include_header=False)
|
||||
|
||||
assert elements[0].text == EXPECTED_TEXT
|
||||
assert all(element.metadata.filename == "test" for element in elements)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_tsv_from_file_path_gets_last_modified_from_filesystem(mocker: MockFixture):
|
||||
filesystem_last_modified = "2024-05-01T15:37:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.tsv.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_tsv(example_doc_path("stanley-cups.tsv"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_tsv_from_file_gets_last_modified_None():
|
||||
with open(example_doc_path("stanley-cups.tsv"), "rb") as f:
|
||||
elements = partition_tsv(file=f)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_tsv_from_file_path_prefers_metadata_last_modified(mocker: MockFixture):
|
||||
filesystem_last_modified = "2024-05-01T15:37:28"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.tsv.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_tsv(
|
||||
example_doc_path("stanley-cups.tsv"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_tsv_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
|
||||
with open(example_doc_path("stanley-cups.tsv"), "rb") as f:
|
||||
elements = partition_tsv(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert elements[0].metadata.last_modified == metadata_last_modified
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", ["stanley-cups.tsv", "stanley-cups-with-emoji.tsv"])
|
||||
def test_partition_tsv_with_json(filename: str):
|
||||
elements = partition_tsv(example_doc_path(filename), include_header=False)
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
# NOTE (jennings) partition_tsv returns a single TableElement per sheet,
|
||||
# so no adding tests for multiple languages like the other partitions
|
||||
def test_partition_tsv_element_metadata_has_languages():
|
||||
filename = "example-docs/stanley-cups-with-emoji.tsv"
|
||||
elements = partition_tsv(filename=filename, include_header=False)
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_tsv_header():
|
||||
elements = partition_tsv(
|
||||
example_doc_path("stanley-cups.tsv"), strategy="fast", include_header=True
|
||||
)
|
||||
|
||||
table = elements[0]
|
||||
assert table.text == "Stanley Cups Unnamed: 1 Unnamed: 2 " + EXPECTED_TEXT_XLSX
|
||||
assert table.metadata.text_as_html is not None
|
||||
assert "<table>" in table.metadata.text_as_html
|
||||
|
||||
|
||||
def test_partition_tsv_supports_chunking_strategy_while_partitioning():
|
||||
elements = partition_tsv(filename=example_doc_path("stanley-cups.tsv"))
|
||||
chunks = chunk_by_title(elements, max_characters=9, combine_text_under_n_chars=0)
|
||||
|
||||
chunk_elements = partition_tsv(
|
||||
example_doc_path("stanley-cups.tsv"),
|
||||
chunking_strategy="by_title",
|
||||
max_characters=9,
|
||||
combine_text_under_n_chars=0,
|
||||
include_header=False,
|
||||
)
|
||||
|
||||
# The same chunks are returned if chunking elements or chunking during partitioning.
|
||||
assert chunk_elements == chunks
|
||||
@@ -0,0 +1,614 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
"""Test-suite for the `unstructured.partition.xlsx` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import sys
|
||||
import tempfile
|
||||
from typing import Any
|
||||
|
||||
import pandas as pd
|
||||
import pandas.testing as pdt
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from test_unstructured.partition.test_constants import (
|
||||
EXPECTED_TABLE_XLSX,
|
||||
EXPECTED_TEXT_XLSX,
|
||||
EXPECTED_TITLE,
|
||||
)
|
||||
from test_unstructured.unit_utils import (
|
||||
FixtureRequest,
|
||||
Mock,
|
||||
assert_round_trips_through_JSON,
|
||||
example_doc_path,
|
||||
function_mock,
|
||||
)
|
||||
from unstructured.cleaners.core import clean_extra_whitespace
|
||||
from unstructured.documents.elements import ListItem, Table, Text, Title
|
||||
from unstructured.errors import UnprocessableEntityError
|
||||
from unstructured.partition.xlsx import (
|
||||
_ConnectedComponent,
|
||||
_SubtableParser,
|
||||
_XlsxPartitionerOptions,
|
||||
partition_xlsx,
|
||||
)
|
||||
|
||||
EXPECTED_FILETYPE = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
|
||||
|
||||
EXCEPTED_PAGE_NAME = "Stanley Cups"
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
# INTEGRATION TESTS
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
# These test `partition_xlsx()` as a whole by calling `partition_xlsx()` and inspecting the
|
||||
# outputs.
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_xlsx_from_filename():
|
||||
elements = partition_xlsx("example-docs/stanley-cups.xlsx", include_header=False)
|
||||
|
||||
assert sum(isinstance(element, Table) for element in elements) == 2
|
||||
assert len(elements) == 4
|
||||
|
||||
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TITLE
|
||||
assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
|
||||
assert elements[1].metadata.text_as_html == EXPECTED_TABLE_XLSX
|
||||
assert elements[1].metadata.page_number == 1
|
||||
assert elements[1].metadata.filetype == EXPECTED_FILETYPE
|
||||
assert elements[1].metadata.page_name == EXCEPTED_PAGE_NAME
|
||||
assert elements[1].metadata.filename == "stanley-cups.xlsx"
|
||||
|
||||
|
||||
def test_partition_xlsx_from_SpooledTemporaryFile_with_emoji():
|
||||
with tempfile.SpooledTemporaryFile() as f:
|
||||
with open("example-docs/emoji.xlsx", "rb") as g:
|
||||
f.write(g.read())
|
||||
|
||||
elements = partition_xlsx(file=f, include_header=False)
|
||||
|
||||
assert sum(isinstance(element, Text) for element in elements) == 1
|
||||
assert len(elements) == 1
|
||||
assert clean_extra_whitespace(elements[0].text) == "🤠😅"
|
||||
|
||||
|
||||
def test_partition_xlsx_raises_on_no_file_or_path_provided():
|
||||
with pytest.raises(ValueError, match="Either 'filename' or 'file' argument must be specif"):
|
||||
partition_xlsx()
|
||||
|
||||
|
||||
def test_partition_xlsx_from_filename_with_metadata_filename():
|
||||
elements = partition_xlsx(
|
||||
"example-docs/stanley-cups.xlsx", metadata_filename="test", include_header=False
|
||||
)
|
||||
|
||||
assert sum(isinstance(element, Table) for element in elements) == 2
|
||||
assert sum(isinstance(element, Title) for element in elements) == 2
|
||||
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TITLE
|
||||
assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
|
||||
assert elements[0].metadata.filename == "test"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("infer_table_structure", [True, False])
|
||||
def test_partition_xlsx_infer_table_structure(infer_table_structure: bool):
|
||||
elements = partition_xlsx(
|
||||
"example-docs/stanley-cups.xlsx", infer_table_structure=infer_table_structure
|
||||
)
|
||||
table_elements = [e for e in elements if isinstance(e, Table)]
|
||||
for table_element in table_elements:
|
||||
table_element_has_text_as_html_field = (
|
||||
hasattr(table_element.metadata, "text_as_html")
|
||||
and table_element.metadata.text_as_html is not None
|
||||
)
|
||||
assert table_element_has_text_as_html_field == infer_table_structure
|
||||
|
||||
|
||||
def test_partition_xlsx_from_filename_with_header():
|
||||
elements = partition_xlsx("example-docs/stanley-cups.xlsx", include_header=True)
|
||||
|
||||
assert len(elements) == 2
|
||||
assert all(isinstance(e, Table) for e in elements)
|
||||
e = elements[0]
|
||||
assert e.text == "Stanley Cups Unnamed: 1 Unnamed: 2 " + EXPECTED_TEXT_XLSX
|
||||
assert e.metadata.text_as_html is not None
|
||||
|
||||
|
||||
def test_partition_xlsx_from_file():
|
||||
with open("example-docs/stanley-cups.xlsx", "rb") as f:
|
||||
elements = partition_xlsx(file=f, include_header=False)
|
||||
|
||||
assert sum(isinstance(element, Table) for element in elements) == 2
|
||||
assert len(elements) == 4
|
||||
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TITLE
|
||||
assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
|
||||
assert elements[1].metadata.text_as_html == EXPECTED_TABLE_XLSX
|
||||
assert elements[1].metadata.page_number == 1
|
||||
assert elements[1].metadata.filetype == EXPECTED_FILETYPE
|
||||
assert elements[1].metadata.page_name == EXCEPTED_PAGE_NAME
|
||||
assert elements[1].metadata.filename is None
|
||||
|
||||
|
||||
def test_partition_xlsx_from_file_like_object_with_name():
|
||||
with open("example-docs/stanley-cups.xlsx", "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
file.name = "stanley-cups-downloaded-from-network.xlsx"
|
||||
|
||||
elements = partition_xlsx(file=file, include_header=False)
|
||||
|
||||
assert sum(isinstance(element, Table) for element in elements) == 2
|
||||
assert len(elements) == 4
|
||||
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TITLE
|
||||
assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
|
||||
assert elements[1].metadata.text_as_html == EXPECTED_TABLE_XLSX
|
||||
assert elements[1].metadata.page_number == 1
|
||||
assert elements[1].metadata.filetype == EXPECTED_FILETYPE
|
||||
assert elements[1].metadata.page_name == EXCEPTED_PAGE_NAME
|
||||
|
||||
|
||||
def test_partition_xlsx_from_file_with_metadata_filename():
|
||||
with open("example-docs/stanley-cups.xlsx", "rb") as f:
|
||||
elements = partition_xlsx(file=f, metadata_filename="test", include_header=False)
|
||||
|
||||
assert sum(isinstance(element, Table) for element in elements) == 2
|
||||
assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
|
||||
assert elements[1].metadata.filename == "test"
|
||||
|
||||
|
||||
def test_partition_xlsx_from_file_with_header():
|
||||
with open("example-docs/stanley-cups.xlsx", "rb") as f:
|
||||
elements = partition_xlsx(file=f, include_header=True)
|
||||
|
||||
assert len(elements) == 2
|
||||
assert all(isinstance(e, Table) for e in elements)
|
||||
e = elements[0]
|
||||
assert e.text == "Stanley Cups Unnamed: 1 Unnamed: 2 " + EXPECTED_TEXT_XLSX
|
||||
assert e.metadata.text_as_html is not None
|
||||
|
||||
|
||||
def test_partition_xlsx_password_protected_raises_exception():
|
||||
with pytest.raises(UnprocessableEntityError):
|
||||
partition_xlsx(filename="example-docs/password_protected.xlsx")
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_xlsx_from_file_path_gets_last_modified_from_filesystem(mocker: MockerFixture):
|
||||
filesystem_last_modified = "2024-05-01T15:37:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.xlsx.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_xlsx(example_doc_path("stanley-cups.xlsx"))
|
||||
|
||||
assert all(e.metadata.last_modified == filesystem_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_xlsx_from_file_gets_last_modified_None():
|
||||
with open(example_doc_path("stanley-cups.xlsx"), "rb") as f:
|
||||
elements = partition_xlsx(file=f)
|
||||
|
||||
assert all(e.metadata.last_modified is None for e in elements)
|
||||
|
||||
|
||||
def test_partition_xlsx_from_file_path_prefers_metadata_last_modified(mocker: MockerFixture):
|
||||
filesystem_last_modified = "2024-05-01T15:37:28"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
mocker.patch(
|
||||
"unstructured.partition.xlsx.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_xlsx(
|
||||
example_doc_path("stanley-cups.xlsx"), metadata_last_modified=metadata_last_modified
|
||||
)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
def test_partition_xlsx_from_file_prefers_metadata_last_modified():
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
|
||||
with open(example_doc_path("stanley-cups.xlsx"), "rb") as f:
|
||||
elements = partition_xlsx(file=f, metadata_last_modified=metadata_last_modified)
|
||||
|
||||
assert all(e.metadata.last_modified == metadata_last_modified for e in elements)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_xlsx_with_json():
|
||||
elements = partition_xlsx(example_doc_path("stanley-cups.xlsx"), include_header=False)
|
||||
assert_round_trips_through_JSON(elements)
|
||||
|
||||
|
||||
def test_partition_xlsx_metadata_language_from_filename():
|
||||
elements = partition_xlsx("example-docs/stanley-cups.xlsx", include_header=False)
|
||||
|
||||
assert sum(isinstance(element, Table) for element in elements) == 2
|
||||
assert len(elements) == 4
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_xlsx_subtables():
|
||||
assert partition_xlsx("example-docs/xlsx-subtable-cases.xlsx") == [
|
||||
Table("a b c d e"),
|
||||
ListItem("f"),
|
||||
Title("a"),
|
||||
Table("b c d e"),
|
||||
Title("a"),
|
||||
Title("b"),
|
||||
Table("c d e f"),
|
||||
Table("a b c d"),
|
||||
ListItem("2. e"),
|
||||
Table("a b c d"),
|
||||
Title("e"),
|
||||
Title("f"),
|
||||
Title("a"),
|
||||
Table("b c d e"),
|
||||
Title("f"),
|
||||
Title("a"),
|
||||
Title("b"),
|
||||
Table("c d e f"),
|
||||
Title("g"),
|
||||
Title("a"),
|
||||
Table("b c d e"),
|
||||
Title("f"),
|
||||
Title("g"),
|
||||
Title("a"),
|
||||
Title("b"),
|
||||
Table("c d e f"),
|
||||
Title("g"),
|
||||
Title("h"),
|
||||
Table("a b c"),
|
||||
Title("a"),
|
||||
Table("b c d"),
|
||||
Table("a b c"),
|
||||
Title("d"),
|
||||
Title("e"),
|
||||
]
|
||||
|
||||
|
||||
def test_partition_xlsx_element_metadata_has_languages():
|
||||
elements = partition_xlsx("example-docs/stanley-cups.xlsx")
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_eml_respects_detect_language_per_element():
|
||||
elements = partition_xlsx(
|
||||
"example-docs/language-docs/eng_spa.xlsx", detect_language_per_element=True
|
||||
)
|
||||
|
||||
langs = {e.metadata.languages[0] for e in elements if e.metadata.languages}
|
||||
assert "eng" in langs
|
||||
assert "spa" in langs
|
||||
|
||||
|
||||
def test_partition_xlsx_with_more_than_1k_cells():
|
||||
old_recursion_limit = sys.getrecursionlimit()
|
||||
try:
|
||||
sys.setrecursionlimit(1000)
|
||||
partition_xlsx("example-docs/more-than-1k-cells.xlsx")
|
||||
finally:
|
||||
sys.setrecursionlimit(old_recursion_limit)
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# OTHER ARGS
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
# -- `find_subtable` -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_xlsx_with_find_subtables_False_emits_one_Table_element_per_worksheet():
|
||||
elements = partition_xlsx("example-docs/stanley-cups.xlsx", find_subtable=False)
|
||||
assert elements == [
|
||||
Table(
|
||||
"Stanley Cups Team Location Stanley Cups Blues STL 1 Flyers PHI 2 Maple Leafs TOR 13"
|
||||
),
|
||||
Table(
|
||||
"Stanley Cups Since 67 Team Location Stanley Cups Blues STL 1 Flyers PHI 2 Maple"
|
||||
" Leafs TOR 0"
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_partition_xlsx_with_find_subtables_False_and_infer_table_structure_False_works():
|
||||
elements = partition_xlsx(
|
||||
"example-docs/stanley-cups.xlsx", find_subtable=False, infer_table_structure=False
|
||||
)
|
||||
assert elements == [
|
||||
Table(
|
||||
"Stanley Cups Team Location Stanley Cups Blues STL 1 Flyers PHI 2 Maple Leafs TOR 13"
|
||||
),
|
||||
Table(
|
||||
"Stanley Cups Since 67 Team Location Stanley Cups Blues STL 1 Flyers PHI 2 Maple"
|
||||
" Leafs TOR 0"
|
||||
),
|
||||
]
|
||||
assert all(e.metadata.text_as_html is None for e in elements)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
# UNIT TESTS
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
# These test components used by `partition_xlsx()` in isolation such that all edge cases can be
|
||||
# exercised.
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
class Describe_XlsxPartitionerOptions:
|
||||
"""Unit-test suite for `unstructured.partition.xlsx._XlsxPartitionerOptions` objects."""
|
||||
|
||||
@pytest.mark.parametrize("arg_value", [True, False])
|
||||
def it_knows_whether_to_find_subtables_within_each_worksheet_or_return_table_per_worksheet(
|
||||
self, arg_value: bool, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["find_subtable"] = arg_value
|
||||
opts = _XlsxPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.find_subtable is arg_value
|
||||
|
||||
@pytest.mark.parametrize(("arg_value", "expected_value"), [(True, 0), (False, None)])
|
||||
def it_knows_the_header_row_index_for_Pandas(
|
||||
self, arg_value: bool, expected_value: int | None, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["include_header"] = arg_value
|
||||
opts = _XlsxPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.header_row_idx == expected_value
|
||||
|
||||
@pytest.mark.parametrize("arg_value", [True, False])
|
||||
def it_knows_whether_to_include_column_headings_in_Table_text_as_html(
|
||||
self, arg_value: bool, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["include_header"] = arg_value
|
||||
opts = _XlsxPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.include_header is arg_value
|
||||
|
||||
@pytest.mark.parametrize("arg_value", [True, False])
|
||||
def it_knows_whether_to_include_text_as_html_in_Table_metadata(
|
||||
self, arg_value: bool, opts_args: dict[str, Any]
|
||||
):
|
||||
opts_args["infer_table_structure"] = arg_value
|
||||
opts = _XlsxPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.infer_table_structure is arg_value
|
||||
|
||||
# -- .last_modified --------------------------------------------------------------------------
|
||||
|
||||
def it_gets_last_modified_from_the_filesystem_when_a_path_is_provided(
|
||||
self, opts_args: dict[str, Any], get_last_modified_date_: Mock
|
||||
):
|
||||
filesystem_last_modified = "2024-04-02T20:32:35"
|
||||
opts_args["file_path"] = "a/b/spreadsheet.xlsx"
|
||||
get_last_modified_date_.return_value = filesystem_last_modified
|
||||
opts = _XlsxPartitionerOptions(**opts_args)
|
||||
|
||||
last_modified = opts.last_modified
|
||||
|
||||
get_last_modified_date_.assert_called_once_with("a/b/spreadsheet.xlsx")
|
||||
assert last_modified == filesystem_last_modified
|
||||
|
||||
def but_it_falls_back_to_None_for_the_last_modified_date_when_no_file_path_is_provided(
|
||||
self, opts_args: dict[str, Any]
|
||||
):
|
||||
file = io.BytesIO(b"abcdefg")
|
||||
opts_args["file"] = file
|
||||
opts = _XlsxPartitionerOptions(**opts_args)
|
||||
|
||||
last_modified = opts.last_modified
|
||||
|
||||
assert last_modified is None
|
||||
|
||||
# -- .metadata_file_path ---------------------------------------------------------------------
|
||||
|
||||
def it_uses_the_file_path_argument_when_provided(self, opts_args: dict[str, Any]):
|
||||
opts_args["file_path"] = "x/y/z.xlsx"
|
||||
opts = _XlsxPartitionerOptions(**opts_args)
|
||||
|
||||
assert opts.metadata_file_path == "x/y/z.xlsx"
|
||||
|
||||
# -- fixtures --------------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture()
|
||||
def get_last_modified_date_(self, request: FixtureRequest):
|
||||
return function_mock(request, "unstructured.partition.xlsx.get_last_modified_date")
|
||||
|
||||
@pytest.fixture()
|
||||
def opts_args(self) -> dict[str, Any]:
|
||||
"""All default arguments for `_XlsxPartitionerOptions`.
|
||||
|
||||
Individual argument values can be changed to suit each test. Makes construction of opts more
|
||||
compact for testing purposes.
|
||||
"""
|
||||
return {
|
||||
"file_path": None,
|
||||
"file": None,
|
||||
"find_subtable": True,
|
||||
"include_header": False,
|
||||
"infer_table_structure": True,
|
||||
}
|
||||
|
||||
|
||||
class Describe_ConnectedComponent:
|
||||
"""Unit-test suite for `unstructured.partition.xlsx._ConnectedComponent` objects."""
|
||||
|
||||
def it_knows_its_top_and_left_extents(self):
|
||||
component = _ConnectedComponent(pd.DataFrame(), {(0, 1), (2, 2), (1, 1), (2, 3), (1, 2)})
|
||||
|
||||
assert component.min_x == 0
|
||||
assert component.max_x == 2
|
||||
|
||||
def it_can_merge_with_another_component_to_make_a_new_component(self):
|
||||
df = pd.DataFrame()
|
||||
component = _ConnectedComponent(df, {(0, 1), (0, 2), (1, 1)})
|
||||
other = _ConnectedComponent(df, {(0, 4), (1, 3), (1, 4)})
|
||||
|
||||
merged = component.merge(other)
|
||||
|
||||
assert merged._worksheet is df
|
||||
assert merged._cell_coordinate_set == {(0, 1), (0, 2), (1, 1), (0, 4), (1, 3), (1, 4)}
|
||||
|
||||
def it_can_extract_the_rectangular_subtable_containing_its_cells_from_the_worksheet(self):
|
||||
worksheet_df = pd.DataFrame(
|
||||
[["a", "b", "c"], [], ["d", "e"], ["f", "g"], [None, "h"], [], ["i"]],
|
||||
index=[0, 1, 2, 3, 4, 5, 6],
|
||||
)
|
||||
cell_coordinate_set = {(2, 0), (2, 1), (3, 0), (3, 1), (4, 1)}
|
||||
component = _ConnectedComponent(worksheet_df, cell_coordinate_set)
|
||||
|
||||
subtable = component.subtable
|
||||
|
||||
print(f"{subtable=}")
|
||||
pdt.assert_frame_equal(
|
||||
subtable, pd.DataFrame([["d", "e"], ["f", "g"], [None, "h"]], index=[2, 3, 4])
|
||||
)
|
||||
|
||||
|
||||
class Describe_SubtableParser:
|
||||
"""Unit-test suite for `unstructured.partition.xlsx._SubtableParser` objects."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("subtable", "expected_value"),
|
||||
[
|
||||
# -- 1. no leading or trailing single-cell rows --
|
||||
(
|
||||
pd.DataFrame([["a", "b"], ["c", "d"]], index=[0, 1]),
|
||||
pd.DataFrame([["a", "b"], ["c", "d"]], index=[0, 1]),
|
||||
),
|
||||
# -- 2. one leading single-cell row --
|
||||
(
|
||||
pd.DataFrame([["a"], ["b", "c"], ["d", "e"]], index=[0, 1, 2]),
|
||||
pd.DataFrame([["b", "c"], ["d", "e"]], index=[1, 2]),
|
||||
),
|
||||
# -- 3. two leading single-cell rows --
|
||||
(
|
||||
pd.DataFrame(
|
||||
[[None, "a"], [None, "b"], ["c", "d"], ["e", "f"]], index=[0, 1, 2, 3]
|
||||
),
|
||||
pd.DataFrame([["c", "d"], ["e", "f"]], index=[2, 3]),
|
||||
),
|
||||
# -- 4. one trailing single-cell row --
|
||||
(
|
||||
pd.DataFrame([["a", "b"], ["c", "d"], [None, "e"]], index=[0, 1, 2]),
|
||||
pd.DataFrame([["a", "b"], ["c", "d"]], index=[0, 1]),
|
||||
),
|
||||
# -- 5. two trailing single-cell rows --
|
||||
(
|
||||
pd.DataFrame([["a", "b"], ["c", "d"], ["e"], ["f"]], index=[0, 1, 2, 3]),
|
||||
pd.DataFrame([["a", "b"], ["c", "d"]], index=[0, 1]),
|
||||
),
|
||||
# -- 6. one leading, one trailing single-cell rows --
|
||||
(
|
||||
pd.DataFrame([["a"], ["b", "c"], ["d", "e"], [None, "f"]], index=[0, 1, 2, 3]),
|
||||
pd.DataFrame([["b", "c"], ["d", "e"]], index=[1, 2]),
|
||||
),
|
||||
# -- 7. two leading, one trailing single-cell rows --
|
||||
(
|
||||
pd.DataFrame([["a"], ["b"], ["c", "d"], ["e", "f"], ["g"]], index=[0, 1, 2, 3, 4]),
|
||||
pd.DataFrame([["c", "d"], ["e", "f"]], index=[2, 3]),
|
||||
),
|
||||
# -- 8. one leading, two trailing single-cell rows --
|
||||
(
|
||||
pd.DataFrame(
|
||||
[[None, "a"], ["b", "c"], ["d", "e"], [None, "f"], [None, "g"]],
|
||||
index=[0, 1, 2, 3, 4],
|
||||
),
|
||||
pd.DataFrame([["b", "c"], ["d", "e"]], index=[1, 2]),
|
||||
),
|
||||
# -- 9. two leading, two trailing single-cell rows --
|
||||
(
|
||||
pd.DataFrame(
|
||||
[["a"], ["b"], ["c", "d"], ["e", "f"], ["g"], ["h"]], index=[0, 1, 2, 3, 4, 5]
|
||||
),
|
||||
pd.DataFrame([["c", "d"], ["e", "f"]], index=[2, 3]),
|
||||
),
|
||||
# -- 10. single-row core-table, no leading or trailing single-cell rows --
|
||||
(
|
||||
pd.DataFrame([["a", "b", "c"]], index=[0]),
|
||||
pd.DataFrame([["a", "b", "c"]], index=[0]),
|
||||
),
|
||||
# -- 11. single-row core-table, one leading single-cell row --
|
||||
(
|
||||
pd.DataFrame([["a"], ["b", "c", "d"]], index=[0, 1]),
|
||||
pd.DataFrame([["b", "c", "d"]], index=[1]),
|
||||
),
|
||||
# -- 12. single-row core-table, two trailing single-cell rows --
|
||||
(
|
||||
pd.DataFrame([["a", "b", "c"], ["d"], ["e"]], index=[0, 1, 2]),
|
||||
pd.DataFrame([["a", "b", "c"]], index=[0]),
|
||||
),
|
||||
],
|
||||
)
|
||||
def it_extracts_the_core_table_from_a_subtable(
|
||||
self, subtable: pd.DataFrame, expected_value: pd.DataFrame
|
||||
):
|
||||
"""core-table is correctly distinguished from leading and trailing single-cell rows."""
|
||||
subtable_parser = _SubtableParser(subtable)
|
||||
|
||||
core_table = subtable_parser.core_table
|
||||
|
||||
assert core_table is not None
|
||||
pdt.assert_frame_equal(core_table, expected_value)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("subtable", "expected_value"),
|
||||
[
|
||||
(pd.DataFrame([["a", "b"], ["c", "d"]]), []),
|
||||
(pd.DataFrame([["a"], ["b", "c"], ["d", "e"]]), ["a"]),
|
||||
(pd.DataFrame([[None, "a"], [None, "b"], ["c", "d"], ["e", "f"]]), ["a", "b"]),
|
||||
(pd.DataFrame([["a", "b"], ["c", "d"], [None, "e"]]), []),
|
||||
(pd.DataFrame([["a", "b"], ["c", "d"], ["e"], ["f"]]), []),
|
||||
(pd.DataFrame([["a"], ["b", "c"], ["d", "e"], [None, "f"]]), ["a"]),
|
||||
(pd.DataFrame([["a"], ["b"], ["c", "d"], ["e", "f"], ["g"]]), ["a", "b"]),
|
||||
(pd.DataFrame([[None, "a"], ["b", "c"], ["d", "e"], [None, "f"], [None, "g"]]), ["a"]),
|
||||
(pd.DataFrame([["a"], ["b"], ["c", "d"], ["e", "f"], ["g"], ["h"]]), ["a", "b"]),
|
||||
(pd.DataFrame([["a", "b", "c"]]), []),
|
||||
(pd.DataFrame([["a"], ["b", "c", "d"]]), ["a"]),
|
||||
(pd.DataFrame([["a", "b", "c"], ["d"], ["e"]]), []),
|
||||
],
|
||||
)
|
||||
def it_extracts_the_leading_single_cell_rows_from_a_subtable(
|
||||
self, subtable: pd.DataFrame, expected_value: pd.DataFrame
|
||||
):
|
||||
subtable_parser = _SubtableParser(subtable)
|
||||
leading_single_cell_row_texts = list(subtable_parser.iter_leading_single_cell_rows_texts())
|
||||
assert leading_single_cell_row_texts == expected_value
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("subtable", "expected_value"),
|
||||
[
|
||||
(pd.DataFrame([["a", "b"], ["c", "d"]]), []),
|
||||
(pd.DataFrame([["a"], ["b", "c"], ["d", "e"]]), []),
|
||||
(pd.DataFrame([[None, "a"], [None, "b"], ["c", "d"], ["e", "f"]]), []),
|
||||
(pd.DataFrame([["a", "b"], ["c", "d"], [None, "e"]]), ["e"]),
|
||||
(pd.DataFrame([["a", "b"], ["c", "d"], ["e"], ["f"]]), ["e", "f"]),
|
||||
(pd.DataFrame([["a"], ["b", "c"], ["d", "e"], [None, "f"]]), ["f"]),
|
||||
(pd.DataFrame([["a"], ["b"], ["c", "d"], ["e", "f"], ["g"]]), ["g"]),
|
||||
(
|
||||
pd.DataFrame([[None, "a"], ["b", "c"], ["d", "e"], [None, "f"], [None, "g"]]),
|
||||
["f", "g"],
|
||||
),
|
||||
(pd.DataFrame([["a"], ["b"], ["c", "d"], ["e", "f"], ["g"], ["h"]]), ["g", "h"]),
|
||||
(pd.DataFrame([["a", "b", "c"]]), []),
|
||||
(pd.DataFrame([["a"], ["b", "c", "d"]]), []),
|
||||
(pd.DataFrame([["a", "b", "c"], ["d"], ["e"]]), ["d", "e"]),
|
||||
],
|
||||
)
|
||||
def it_extracts_the_trailing_single_cell_rows_from_a_subtable(
|
||||
self, subtable: pd.DataFrame, expected_value: pd.DataFrame
|
||||
):
|
||||
subtable_parser = _SubtableParser(subtable)
|
||||
|
||||
trailing_single_cell_row_texts = list(
|
||||
subtable_parser.iter_trailing_single_cell_rows_texts()
|
||||
)
|
||||
|
||||
assert trailing_single_cell_row_texts == expected_value
|
||||
@@ -0,0 +1,229 @@
|
||||
"""Test-suite for `unstructured.partition.xml` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from test_unstructured.unit_utils import example_doc_path
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import NarrativeText, Title
|
||||
from unstructured.partition.json import partition_json
|
||||
from unstructured.partition.utils.constants import UNSTRUCTURED_INCLUDE_DEBUG_METADATA
|
||||
from unstructured.partition.xml import partition_xml
|
||||
from unstructured.staging.base import elements_to_json
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", ["factbook.xml", "factbook-utf-16.xml"])
|
||||
def test_partition_xml_from_filename(filename: str):
|
||||
file_path = example_doc_path(filename)
|
||||
elements = partition_xml(filename=file_path, xml_keep_tags=False)
|
||||
|
||||
assert elements[0].text == "United States"
|
||||
assert elements[0].metadata.filename == filename
|
||||
if UNSTRUCTURED_INCLUDE_DEBUG_METADATA:
|
||||
assert {element.metadata.detection_origin for element in elements} == {"xml"}
|
||||
|
||||
|
||||
def test_partition_xml_from_filename_with_metadata_filename():
|
||||
elements = partition_xml(
|
||||
example_doc_path("factbook.xml"), xml_keep_tags=False, metadata_filename="test"
|
||||
)
|
||||
|
||||
assert elements[0].text == "United States"
|
||||
assert elements[0].metadata.filename == "test"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", ["factbook.xml", "factbook-utf-16.xml"])
|
||||
def test_partition_xml_from_file(filename: str):
|
||||
file_path = example_doc_path(filename)
|
||||
with open(file_path, "rb") as f:
|
||||
elements = partition_xml(file=f, xml_keep_tags=False, metadata_filename=file_path)
|
||||
|
||||
assert elements[0].text == "United States"
|
||||
assert elements[0].metadata.filename == filename
|
||||
|
||||
|
||||
def test_partition_xml_from_file_with_metadata_filename():
|
||||
with open(example_doc_path("factbook.xml"), "rb") as f:
|
||||
elements = partition_xml(file=f, xml_keep_tags=False, metadata_filename="test")
|
||||
|
||||
assert elements[0].text == "United States"
|
||||
assert elements[0].metadata.filename == "test"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", ["factbook.xml", "factbook-utf-16.xml"])
|
||||
def test_partition_xml_from_file_rb(filename: str):
|
||||
file_path = example_doc_path(filename)
|
||||
with open(file_path, "rb") as f:
|
||||
elements = partition_xml(file=f, xml_keep_tags=False, metadata_filename=file_path)
|
||||
|
||||
assert elements[0].text == "United States"
|
||||
assert elements[0].metadata.filename == filename
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", ["factbook.xml", "factbook-utf-16.xml"])
|
||||
def test_partition_xml_from_filename_with_tags_default_encoding(filename: str):
|
||||
file_path = example_doc_path(filename)
|
||||
elements = partition_xml(filename=file_path, xml_keep_tags=True)
|
||||
|
||||
assert "<leader>Joe Biden</leader>" in elements[0].text
|
||||
assert elements[0].metadata.filename == filename
|
||||
|
||||
|
||||
def test_partition_xml_from_text_with_tags():
|
||||
with open(example_doc_path("factbook.xml")) as f:
|
||||
text = f.read()
|
||||
elements = partition_xml(text=text, xml_keep_tags=True)
|
||||
|
||||
assert "<leader>Joe Biden</leader>" in elements[0].text
|
||||
|
||||
|
||||
def test_partition_xml_from_filename_with_tags_raises_encoding_error():
|
||||
with pytest.raises(UnicodeDecodeError):
|
||||
partition_xml(example_doc_path("factbook-utf-16.xml"), xml_keep_tags=True, encoding="utf-8")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", ["factbook.xml", "factbook-utf-16.xml"])
|
||||
def test_partition_xml_from_file_with_tags_default_encoding(filename: str):
|
||||
file_path = example_doc_path(filename)
|
||||
with open(file_path, "rb") as f:
|
||||
elements = partition_xml(file=f, xml_keep_tags=True, metadata_filename=file_path)
|
||||
|
||||
assert "<leader>Joe Biden</leader>" in elements[0].text
|
||||
assert elements[0].metadata.filename == filename
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", ["factbook.xml", "factbook-utf-16.xml"])
|
||||
def test_partition_xml_from_file_rb_with_tags_default_encoding(filename: str):
|
||||
file_path = example_doc_path(filename)
|
||||
with open(file_path, "rb") as f:
|
||||
elements = partition_xml(file=f, xml_keep_tags=True, metadata_filename=file_path)
|
||||
|
||||
assert "<leader>Joe Biden</leader>" in elements[0].text
|
||||
assert elements[0].metadata.filename == filename
|
||||
|
||||
|
||||
def test_partition_xml_from_file_rb_with_tags_raises_encoding_error():
|
||||
with pytest.raises(UnicodeDecodeError):
|
||||
with open(example_doc_path("factbook-utf-16.xml"), "rb") as f:
|
||||
partition_xml(
|
||||
file=f,
|
||||
xml_keep_tags=True,
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
# -- .metadata.filetype --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_xml_gets_the_XML_mime_type_in_metadata_filetype():
|
||||
XML_MIME_TYPE = "application/xml"
|
||||
elements = partition_xml(example_doc_path("factbook.xml"))
|
||||
assert all(e.metadata.filetype == XML_MIME_TYPE for e in elements), (
|
||||
f"Expected all elements to have '{XML_MIME_TYPE}' as their filetype, but got:"
|
||||
f" {repr(elements[0].metadata.filetype)}"
|
||||
)
|
||||
|
||||
|
||||
# -- .metadata.last_modified ---------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_partition_xml_from_file_path_gets_last_modified_from_filesystem(mocker: MockerFixture):
|
||||
mocked_last_modification_date = "2029-07-05T09:24:28"
|
||||
|
||||
mocker.patch(
|
||||
"unstructured.partition.xml.get_last_modified_date",
|
||||
return_value=mocked_last_modification_date,
|
||||
)
|
||||
|
||||
elements = partition_xml(filename="example-docs/factbook.xml")
|
||||
|
||||
assert elements[0].metadata.last_modified == mocked_last_modification_date
|
||||
|
||||
|
||||
def test_partition_xml_from_file_gets_last_modified_None():
|
||||
with open("example-docs/factbook.xml", "rb") as f:
|
||||
elements = partition_xml(file=f)
|
||||
|
||||
assert elements[0].metadata.last_modified is None
|
||||
|
||||
|
||||
def test_partition_xml_from_file_path_prefers_metadata_last_modified(mocker: MockerFixture):
|
||||
filesystem_last_modified = "2029-07-05T09:24:28"
|
||||
metadata_last_modified = "2020-07-05T09:24:28"
|
||||
|
||||
mocker.patch(
|
||||
"unstructured.partition.xml.get_last_modified_date", return_value=filesystem_last_modified
|
||||
)
|
||||
|
||||
elements = partition_xml(
|
||||
filename="example-docs/factbook.xml",
|
||||
metadata_last_modified=metadata_last_modified,
|
||||
)
|
||||
|
||||
assert elements[0].metadata.last_modified == metadata_last_modified
|
||||
|
||||
|
||||
def test_partition_xml_from_file_prefers_metadata_last_modified():
|
||||
with open("example-docs/factbook.xml", "rb") as f:
|
||||
elements = partition_xml(file=f, metadata_last_modified="2029-07-05T09:24:28")
|
||||
|
||||
assert elements[0].metadata.last_modified == "2029-07-05T09:24:28"
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", ["factbook.xml", "factbook-utf-16.xml"])
|
||||
def test_partition_xml_with_json(filename: str):
|
||||
file_path = example_doc_path(filename)
|
||||
elements = partition_xml(filename=file_path, xml_keep_tags=False)
|
||||
test_elements = partition_json(text=elements_to_json(elements))
|
||||
|
||||
assert len(elements) == len(test_elements)
|
||||
assert elements[0].metadata.page_number == test_elements[0].metadata.page_number
|
||||
assert elements[0].metadata.filename == test_elements[0].metadata.filename
|
||||
|
||||
for i in range(len(elements)):
|
||||
assert elements[i] == test_elements[i]
|
||||
|
||||
|
||||
def test_partition_xml_with_narrative_line_breaks():
|
||||
xml_text = """<xml>
|
||||
<parrot>
|
||||
<name>Conure</name>
|
||||
<description>A conure is a very friendly bird.
|
||||
Conures are feathery and like to dance.
|
||||
</description>
|
||||
</parrot>
|
||||
</xml>"""
|
||||
|
||||
elements = partition_xml(text=xml_text)
|
||||
assert elements[0] == Title("Conure")
|
||||
assert isinstance(elements[1], NarrativeText)
|
||||
assert str(elements[1]).startswith("A conure is a very friendly bird.")
|
||||
assert str(elements[1]).strip().endswith("Conures are feathery and like to dance.")
|
||||
|
||||
|
||||
def test_add_chunking_strategy_on_partition_xml():
|
||||
file_path = example_doc_path("factbook.xml")
|
||||
elements = partition_xml(file_path)
|
||||
chunk_elements = partition_xml(file_path, chunking_strategy="by_title")
|
||||
chunks = chunk_by_title(elements)
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_partition_xml_element_metadata_has_languages():
|
||||
file_path = example_doc_path("factbook.xml")
|
||||
elements = partition_xml(file_path)
|
||||
assert elements[0].metadata.languages == ["eng"]
|
||||
|
||||
|
||||
def test_partition_xml_respects_detect_language_per_element():
|
||||
elements = partition_xml(
|
||||
example_doc_path("language-docs/eng_spa_mult.xml"), detect_language_per_element=True
|
||||
)
|
||||
langs = [element.metadata.languages for element in elements]
|
||||
assert langs == [["eng"], ["spa", "eng"], ["eng"], ["eng"], ["spa"]]
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user