修改为东南天坐标系

This commit is contained in:
2026-01-20 09:49:52 +08:00
parent 9538757047
commit 333fad40ac
7201 changed files with 1030888 additions and 85410 deletions

View File

@@ -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"

View File

@@ -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,
)

View File

@@ -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))

View File

@@ -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)

View File

@@ -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)

View File

@@ -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&gt;11,442,231&lt;/td&gt;&lt;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&quot;="">
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__

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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"

View File

@@ -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",
),
]

View File

@@ -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

View File

@@ -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}

View File

@@ -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,
),
]
)

View File

@@ -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

View File

@@ -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,
}

View File

@@ -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

View File

@@ -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)

View File

@@ -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

View File

@@ -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
]

View File

@@ -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>"
)

View File

@@ -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")

View File

@@ -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"),
]

View File

@@ -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")

View File

@@ -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"]]

View File

@@ -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)

View File

@@ -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

View File

@@ -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": {},
}

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

View File

@@ -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

View File

@@ -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")

View File

@@ -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)

View File

@@ -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"]]

View File

@@ -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

View File

@@ -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

View File

@@ -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 REGISTRANTS COMMON EQUITY, RELATED STOCKHOLDER MATTERS AND "
"ISSUER PURCHASES OF EQUITY SECURITIES",
False,
),
(
"Item 5(a).: Market For Registrants Common Equity, Related Stockholder Matters and "
"Issuer Purchases of Equity Securities",
False,
),
(
"There is a market for registrants 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(" ")

View File

@@ -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

View File

@@ -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

View File

@@ -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