mirror of
https://github.com/Unstructured-IO/unstructured.git
synced 2025-08-19 06:09:32 +00:00
db4c5dfdf7
Added the CoordinateSystem class for tracking the system in which coordinates are represented, and changing the system if desired.
800 lines
28 KiB
Python
800 lines
28 KiB
Python
import json
|
||
import os
|
||
import pathlib
|
||
import warnings
|
||
from importlib import import_module
|
||
from unittest.mock import patch
|
||
|
||
import docx
|
||
import pytest
|
||
|
||
from test_unstructured.partition.test_constants import EXPECTED_TABLE, EXPECTED_TEXT
|
||
from unstructured.cleaners.core import clean_extra_whitespace
|
||
from unstructured.documents.elements import (
|
||
Address,
|
||
ElementMetadata,
|
||
ListItem,
|
||
NarrativeText,
|
||
PageBreak,
|
||
Table,
|
||
Text,
|
||
Title,
|
||
)
|
||
from unstructured.file_utils.filetype import FILETYPE_TO_MIMETYPE, FileType
|
||
from unstructured.partition import auto
|
||
from unstructured.partition.auto import partition
|
||
from unstructured.partition.common import convert_office_doc
|
||
from unstructured.staging.base import elements_to_json
|
||
|
||
DIRECTORY = pathlib.Path(__file__).parent.resolve()
|
||
EXAMPLE_DOCS_DIRECTORY = os.path.join(DIRECTORY, "..", "..", "example-docs")
|
||
|
||
EXPECTED_EMAIL_OUTPUT = [
|
||
NarrativeText(text="This is a test email to use for unit tests."),
|
||
Title(text="Important points:"),
|
||
ListItem(text="Roses are red"),
|
||
ListItem(text="Violets are blue"),
|
||
]
|
||
|
||
EML_TEST_FILE = "eml/fake-email.eml"
|
||
|
||
is_in_docker = os.path.exists("/.dockerenv")
|
||
|
||
|
||
def test_auto_partition_email_from_filename():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, EML_TEST_FILE)
|
||
elements = partition(filename=filename, strategy="hi_res")
|
||
assert len(elements) > 0
|
||
assert elements == EXPECTED_EMAIL_OUTPUT
|
||
assert elements[0].metadata.filename == os.path.basename(filename)
|
||
assert elements[0].metadata.file_directory == os.path.split(filename)[0]
|
||
|
||
|
||
def test_auto_partition_email_from_file():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, EML_TEST_FILE)
|
||
with open(filename) as f:
|
||
elements = partition(file=f, strategy="hi_res")
|
||
assert len(elements) > 0
|
||
assert elements == EXPECTED_EMAIL_OUTPUT
|
||
|
||
|
||
def test_auto_partition_email_from_file_rb():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, EML_TEST_FILE)
|
||
with open(filename, "rb") as f:
|
||
elements = partition(file=f, strategy="hi_res")
|
||
assert len(elements) > 0
|
||
assert elements == EXPECTED_EMAIL_OUTPUT
|
||
|
||
|
||
@pytest.fixture()
|
||
def mock_docx_document():
|
||
document = docx.Document()
|
||
|
||
document.add_paragraph("These are a few of my favorite things:", style="Heading 1")
|
||
# NOTE(robinson) - this should get picked up as a list item due to the •
|
||
document.add_paragraph("• Parrots", style="Normal")
|
||
document.add_paragraph("Hockey", style="List Bullet")
|
||
# NOTE(robinson) - this should get picked up as a title
|
||
document.add_paragraph("Analysis", style="Normal")
|
||
# NOTE(robinson) - this should get dropped because it is empty
|
||
document.add_paragraph("", style="Normal")
|
||
# NOTE(robinson) - this should get picked up as a narrative text
|
||
document.add_paragraph("This is my first thought. This is my second thought.", style="Normal")
|
||
document.add_paragraph("This is my third thought.", style="Body Text")
|
||
# NOTE(robinson) - this should just be regular text
|
||
document.add_paragraph("2023")
|
||
|
||
return document
|
||
|
||
|
||
@pytest.fixture()
|
||
def expected_docx_elements():
|
||
return [
|
||
Title("These are a few of my favorite things:"),
|
||
ListItem("Parrots"),
|
||
ListItem("Hockey"),
|
||
Title("Analysis"),
|
||
NarrativeText("This is my first thought. This is my second thought."),
|
||
NarrativeText("This is my third thought."),
|
||
Text("2023"),
|
||
]
|
||
|
||
|
||
def test_auto_partition_docx_with_filename(mock_docx_document, expected_docx_elements, tmpdir):
|
||
filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
||
mock_docx_document.save(filename)
|
||
|
||
elements = partition(filename=filename, strategy="hi_res")
|
||
assert elements == expected_docx_elements
|
||
assert elements[0].metadata.filename == os.path.basename(filename)
|
||
|
||
|
||
def test_auto_partition_docx_with_file(mock_docx_document, expected_docx_elements, tmpdir):
|
||
filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
||
mock_docx_document.save(filename)
|
||
|
||
with open(filename, "rb") as f:
|
||
elements = partition(file=f, strategy="hi_res")
|
||
assert elements == expected_docx_elements
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("pass_file_filename", "content_type"),
|
||
[(False, None), (False, "application/msword"), (True, "application/msword"), (True, None)],
|
||
)
|
||
def test_auto_partition_doc_with_filename(
|
||
mock_docx_document,
|
||
expected_docx_elements,
|
||
tmpdir,
|
||
pass_file_filename,
|
||
content_type,
|
||
):
|
||
docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
||
doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
|
||
mock_docx_document.save(docx_filename)
|
||
convert_office_doc(docx_filename, tmpdir.dirname, "doc")
|
||
file_filename = doc_filename if pass_file_filename else None
|
||
elements = partition(
|
||
filename=doc_filename,
|
||
file_filename=file_filename,
|
||
content_type=content_type,
|
||
strategy="hi_res",
|
||
)
|
||
assert elements == expected_docx_elements
|
||
assert elements[0].metadata.filename == "mock_document.doc"
|
||
assert elements[0].metadata.file_directory == tmpdir.dirname
|
||
|
||
|
||
# NOTE(robinson) - the application/x-ole-storage mime type is not specific enough to
|
||
# determine that the file is an .doc document
|
||
@pytest.mark.xfail()
|
||
def test_auto_partition_doc_with_file(mock_docx_document, expected_docx_elements, tmpdir):
|
||
docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
||
doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
|
||
mock_docx_document.save(docx_filename)
|
||
convert_office_doc(docx_filename, tmpdir.dirname, "doc")
|
||
|
||
with open(doc_filename, "rb") as f:
|
||
elements = partition(file=f, strategy="hi_res")
|
||
assert elements == expected_docx_elements
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("pass_file_filename", "content_type"),
|
||
[(False, None), (False, "text/html"), (True, "text/html"), (True, None)],
|
||
)
|
||
def test_auto_partition_html_from_filename(pass_file_filename, content_type):
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "example-10k.html")
|
||
file_filename = filename if pass_file_filename else None
|
||
elements = partition(
|
||
filename=filename,
|
||
file_filename=file_filename,
|
||
content_type=content_type,
|
||
strategy="hi_res",
|
||
)
|
||
assert len(elements) > 0
|
||
assert elements[0].metadata.filename == os.path.basename(filename)
|
||
assert elements[0].metadata.file_directory == os.path.split(filename)[0]
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("pass_file_filename", "content_type"),
|
||
[(False, None), (False, "text/html"), (True, "text/html"), (True, None)],
|
||
)
|
||
def test_auto_partition_html_from_file(pass_file_filename, content_type):
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-html.html")
|
||
file_filename = filename if pass_file_filename else None
|
||
with open(filename) as f:
|
||
elements = partition(
|
||
file=f,
|
||
file_filename=file_filename,
|
||
content_type=content_type,
|
||
strategy="hi_res",
|
||
)
|
||
assert len(elements) > 0
|
||
|
||
|
||
def test_auto_partition_html_from_file_rb():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-html.html")
|
||
with open(filename, "rb") as f:
|
||
elements = partition(file=f, strategy="hi_res")
|
||
assert len(elements) > 0
|
||
|
||
|
||
def test_auto_partition_json_from_filename():
|
||
"""Test auto-processing an unstructured json output file by filename."""
|
||
filename = os.path.join(
|
||
EXAMPLE_DOCS_DIRECTORY,
|
||
"..",
|
||
"test_unstructured_ingest",
|
||
"expected-structured-output",
|
||
"azure-blob-storage",
|
||
"spring-weather.html.json",
|
||
)
|
||
with open(filename) as json_f:
|
||
json_data = json.load(json_f)
|
||
json_elems = json.loads(elements_to_json(partition(filename=filename, strategy="hi_res")))
|
||
for elem in json_elems:
|
||
# coordinates are always in the element data structures, even if None
|
||
elem.pop("coordinates")
|
||
elem.pop("coordinate_system")
|
||
elem.pop("layout_width")
|
||
elem.pop("layout_height")
|
||
elem.pop("metadata")
|
||
for elem in json_data:
|
||
elem.pop("metadata")
|
||
assert json_data == json_elems
|
||
|
||
|
||
@pytest.mark.xfail(
|
||
reason="parsed as text not json, https://github.com/Unstructured-IO/unstructured/issues/492",
|
||
)
|
||
def test_auto_partition_json_from_file():
|
||
"""Test auto-processing an unstructured json output file by file handle."""
|
||
filename = os.path.join(
|
||
EXAMPLE_DOCS_DIRECTORY,
|
||
"..",
|
||
"test_unstructured_ingest",
|
||
"expected-structured-output",
|
||
"azure-blob-storage",
|
||
"spring-weather.html.json",
|
||
)
|
||
with open(filename) as json_f:
|
||
json_data = json.load(json_f)
|
||
with open(filename, encoding="utf-8") as partition_f:
|
||
json_elems = json.loads(elements_to_json(partition(file=partition_f, strategy="hi_res")))
|
||
for elem in json_elems:
|
||
# coordinates are always in the element data structures, even if None
|
||
elem.pop("coordinates")
|
||
elem.pop("coordinate_system")
|
||
assert json_data == json_elems
|
||
|
||
|
||
EXPECTED_TEXT_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"),
|
||
]
|
||
|
||
|
||
def test_auto_partition_text_from_filename():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-text.txt")
|
||
elements = partition(filename=filename, strategy="hi_res")
|
||
assert len(elements) > 0
|
||
assert elements == EXPECTED_TEXT_OUTPUT
|
||
assert elements[0].metadata.filename == os.path.basename(filename)
|
||
assert elements[0].metadata.file_directory == os.path.split(filename)[0]
|
||
|
||
|
||
def test_auto_partition_text_from_file():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-text.txt")
|
||
with open(filename) as f:
|
||
elements = partition(file=f, strategy="hi_res")
|
||
assert len(elements) > 0
|
||
assert elements == EXPECTED_TEXT_OUTPUT
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("pass_file_filename", "content_type"),
|
||
[(False, None), (False, "application/pdf"), (True, "application/pdf"), (True, None)],
|
||
)
|
||
def test_auto_partition_pdf_from_filename(pass_file_filename, content_type, request):
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.pdf")
|
||
file_filename = filename if pass_file_filename else None
|
||
|
||
elements = partition(
|
||
filename=filename,
|
||
file_filename=file_filename,
|
||
content_type=content_type,
|
||
strategy="hi_res",
|
||
)
|
||
|
||
assert isinstance(elements[0], Title)
|
||
assert elements[0].text.startswith("LayoutParser")
|
||
|
||
assert elements[0].metadata.filename == os.path.basename(filename)
|
||
assert elements[0].metadata.file_directory == os.path.split(filename)[0]
|
||
|
||
# NOTE(alan): Xfail since new model skips the word Zejiang
|
||
request.applymarker(pytest.mark.xfail)
|
||
|
||
assert isinstance(elements[1], NarrativeText)
|
||
assert elements[1].text.startswith("Zejiang Shen")
|
||
|
||
|
||
def test_auto_partition_pdf_uses_table_extraction():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.pdf")
|
||
with patch(
|
||
"unstructured_inference.inference.layout.process_file_with_model",
|
||
) as mock_process_file_with_model:
|
||
partition(filename, pdf_infer_table_structure=True, strategy="hi_res")
|
||
assert mock_process_file_with_model.call_args[1]["extract_tables"]
|
||
|
||
|
||
def test_auto_partition_pdf_with_fast_strategy():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.pdf")
|
||
|
||
mock_return = [NarrativeText("Hello there!")]
|
||
with patch.object(auto, "partition_pdf", return_value=mock_return) as mock_partition:
|
||
partition(filename=filename, strategy="fast")
|
||
|
||
mock_partition.assert_called_once_with(
|
||
filename=filename,
|
||
file=None,
|
||
url=None,
|
||
include_page_breaks=False,
|
||
infer_table_structure=False,
|
||
strategy="fast",
|
||
ocr_languages="eng",
|
||
)
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("pass_file_filename", "content_type"),
|
||
[(False, None), (False, "application/pdf"), (True, "application/pdf"), (True, None)],
|
||
)
|
||
def test_auto_partition_pdf_from_file(pass_file_filename, content_type, request):
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.pdf")
|
||
file_filename = filename if pass_file_filename else None
|
||
|
||
with open(filename, "rb") as f:
|
||
elements = partition(
|
||
file=f,
|
||
file_filename=file_filename,
|
||
content_type=content_type,
|
||
strategy="hi_res",
|
||
)
|
||
|
||
assert isinstance(elements[0], Title)
|
||
assert elements[0].text.startswith("LayoutParser")
|
||
|
||
# NOTE(alan): Xfail since new model misses the first word Zejiang
|
||
request.applymarker(pytest.mark.xfail)
|
||
|
||
assert isinstance(elements[1], NarrativeText)
|
||
assert elements[1].text.startswith("Zejiang Shen")
|
||
|
||
|
||
def test_partition_pdf_doesnt_raise_warning():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.pdf")
|
||
# NOTE(robinson): This is the recommended way to check that no warning is emitted,
|
||
# per the pytest docs.
|
||
# ref: https://docs.pytest.org/en/7.0.x/how-to/capture-warnings.html
|
||
# #additional-use-cases-of-warnings-in-tests
|
||
with warnings.catch_warnings():
|
||
warnings.simplefilter("error")
|
||
partition(filename=filename, strategy="hi_res")
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("pass_file_filename", "content_type"),
|
||
[(False, None), (False, "image/jpeg"), (True, "image/jpeg"), (True, None)],
|
||
)
|
||
def test_auto_partition_jpg(pass_file_filename, content_type):
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.jpg")
|
||
file_filename = filename if pass_file_filename else None
|
||
elements = partition(
|
||
filename=filename,
|
||
file_filename=file_filename,
|
||
content_type=content_type,
|
||
strategy="auto",
|
||
)
|
||
assert len(elements) > 0
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("pass_file_filename", "content_type"),
|
||
[(False, None), (False, "image/jpeg"), (True, "image/jpeg"), (True, None)],
|
||
)
|
||
def test_auto_partition_jpg_from_file(pass_file_filename, content_type):
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.jpg")
|
||
file_filename = filename if pass_file_filename else None
|
||
with open(filename, "rb") as f:
|
||
elements = partition(
|
||
file=f,
|
||
file_filename=file_filename,
|
||
content_type=content_type,
|
||
strategy="auto",
|
||
)
|
||
assert len(elements) > 0
|
||
|
||
|
||
def test_auto_partition_raises_with_bad_type(monkeypatch):
|
||
monkeypatch.setattr(auto, "detect_filetype", lambda *args, **kwargs: None)
|
||
with pytest.raises(ValueError):
|
||
partition(filename="made-up.fake", strategy="hi_res")
|
||
|
||
|
||
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!"),
|
||
]
|
||
|
||
|
||
def test_auto_partition_pptx_from_filename():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-power-point.pptx")
|
||
elements = partition(filename=filename, strategy="hi_res")
|
||
assert elements == EXPECTED_PPTX_OUTPUT
|
||
assert elements[0].metadata.filename == os.path.basename(filename)
|
||
assert elements[0].metadata.file_directory == os.path.split(filename)[0]
|
||
|
||
|
||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||
def test_auto_partition_ppt_from_filename():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-power-point.ppt")
|
||
elements = partition(filename=filename, strategy="hi_res")
|
||
assert elements == EXPECTED_PPTX_OUTPUT
|
||
assert elements[0].metadata.filename == os.path.basename(filename)
|
||
assert elements[0].metadata.file_directory == os.path.split(filename)[0]
|
||
|
||
|
||
def test_auto_with_page_breaks():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.pdf")
|
||
elements = partition(filename=filename, include_page_breaks=True, strategy="hi_res")
|
||
assert any(isinstance(element, PageBreak) for element in elements)
|
||
|
||
|
||
def test_auto_partition_epub_from_filename():
|
||
filename = os.path.join(DIRECTORY, "..", "..", "example-docs", "winter-sports.epub")
|
||
elements = partition(filename=filename, strategy="hi_res")
|
||
assert len(elements) > 0
|
||
assert elements[0].text.startswith("The Project Gutenberg eBook of Winter Sports")
|
||
|
||
|
||
def test_auto_partition_epub_from_file():
|
||
filename = os.path.join(DIRECTORY, "..", "..", "example-docs", "winter-sports.epub")
|
||
with open(filename, "rb") as f:
|
||
elements = partition(file=f, strategy="hi_res")
|
||
assert len(elements) > 0
|
||
assert elements[0].text.startswith("The Project Gutenberg eBook of Winter Sports")
|
||
|
||
|
||
EXPECTED_MSG_OUTPUT = [
|
||
NarrativeText(text="This is a test email to use for unit tests."),
|
||
Title(text="Important points:"),
|
||
ListItem(text="Roses are red"),
|
||
ListItem(text="Violets are blue"),
|
||
]
|
||
|
||
|
||
def test_auto_partition_msg_from_filename():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-email.msg")
|
||
elements = partition(filename=filename, strategy="hi_res")
|
||
assert elements == EXPECTED_MSG_OUTPUT
|
||
|
||
|
||
def test_auto_partition_rtf_from_filename():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-doc.rtf")
|
||
elements = partition(filename=filename, strategy="hi_res")
|
||
assert elements[0] == Title("My First Heading")
|
||
|
||
|
||
def test_auto_partition_from_url():
|
||
url = "https://raw.githubusercontent.com/Unstructured-IO/unstructured/main/LICENSE.md"
|
||
elements = partition(url=url, content_type="text/plain", strategy="hi_res")
|
||
assert elements[0] == Title("Apache License")
|
||
assert elements[0].metadata.url == url
|
||
|
||
|
||
def test_partition_md_works_with_embedded_html():
|
||
url = "https://raw.githubusercontent.com/Unstructured-IO/unstructured/main/README.md"
|
||
elements = partition(url=url, content_type="text/markdown", strategy="hi_res")
|
||
elements[0].text
|
||
unstructured_found = False
|
||
for element in elements:
|
||
if "unstructured" in elements[0].text:
|
||
unstructured_found = True
|
||
break
|
||
assert unstructured_found is True
|
||
|
||
|
||
def test_auto_partition_warns_if_header_set_and_not_url(caplog):
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, EML_TEST_FILE)
|
||
partition(filename=filename, headers={"Accept": "application/pdf"}, strategy="hi_res")
|
||
assert caplog.records[0].levelname == "WARNING"
|
||
|
||
|
||
def test_auto_partition_works_with_unstructured_jsons():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "spring-weather.html.json")
|
||
elements = partition(filename=filename, strategy="hi_res")
|
||
assert elements[0].text == "News Around NOAA"
|
||
|
||
|
||
def test_auto_partition_works_with_unstructured_jsons_from_file():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "spring-weather.html.json")
|
||
with open(filename, "rb") as f:
|
||
elements = partition(file=f, strategy="hi_res")
|
||
assert elements[0].text == "News Around NOAA"
|
||
|
||
|
||
def test_auto_partition_odt_from_filename():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake.odt")
|
||
elements = partition(filename=filename, strategy="hi_res")
|
||
assert elements == [Title("Lorem ipsum dolor sit amet.")]
|
||
|
||
|
||
def test_auto_partition_odt_from_file():
|
||
filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake.odt")
|
||
with open(filename, "rb") as f:
|
||
elements = partition(file=f, strategy="hi_res")
|
||
|
||
assert elements == [Title("Lorem ipsum dolor sit amet.")]
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("content_type", "routing_func", "expected"),
|
||
[
|
||
("application/json", "json", "application/json"),
|
||
("text/html", "html", "text/html"),
|
||
("jdsfjdfsjkds", "pdf", None),
|
||
],
|
||
)
|
||
def test_auto_adds_filetype_to_metadata(content_type, routing_func, expected):
|
||
with patch(
|
||
f"unstructured.partition.auto.partition_{routing_func}",
|
||
lambda *args, **kwargs: [Text("text 1"), Text("text 2")],
|
||
):
|
||
elements = partition("example-docs/layout-parser-paper-fast.pdf", content_type=content_type)
|
||
assert len(elements) == 2
|
||
assert all(el.metadata.filetype == expected for el in elements)
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("content_type", "expected"),
|
||
[
|
||
("application/pdf", FILETYPE_TO_MIMETYPE[FileType.PDF]),
|
||
(None, FILETYPE_TO_MIMETYPE[FileType.PDF]),
|
||
],
|
||
)
|
||
def test_auto_filetype_overrides_file_specific(content_type, expected):
|
||
pdf_metadata = ElementMetadata(filetype="imapdf")
|
||
with patch(
|
||
"unstructured.partition.auto.partition_pdf",
|
||
lambda *args, **kwargs: [
|
||
Text("text 1", metadata=pdf_metadata),
|
||
Text("text 2", metadata=pdf_metadata),
|
||
],
|
||
):
|
||
elements = partition("example-docs/layout-parser-paper-fast.pdf", content_type=content_type)
|
||
assert len(elements) == 2
|
||
assert all(el.metadata.filetype == expected for el in elements)
|
||
|
||
|
||
supported_filetypes = [
|
||
_
|
||
for _ in FileType
|
||
if _
|
||
not in (
|
||
FileType.UNK,
|
||
FileType.ZIP,
|
||
FileType.XLS,
|
||
)
|
||
]
|
||
|
||
|
||
FILETYPE_TO_MODULE = {
|
||
FileType.JPG: "image",
|
||
FileType.PNG: "image",
|
||
FileType.TXT: "text",
|
||
FileType.EML: "email",
|
||
}
|
||
|
||
|
||
@pytest.mark.parametrize("filetype", supported_filetypes)
|
||
def test_file_specific_produces_correct_filetype(filetype: FileType):
|
||
if filetype in (FileType.JPG, FileType.PNG, FileType.EMPTY):
|
||
pytest.skip()
|
||
extension = filetype.name.lower()
|
||
filetype_module = (
|
||
extension if filetype not in FILETYPE_TO_MODULE else FILETYPE_TO_MODULE[filetype]
|
||
)
|
||
fun_name = "partition_" + filetype_module
|
||
module = import_module(f"unstructured.partition.{filetype_module}") # noqa
|
||
fun = eval(f"module.{fun_name}")
|
||
for file in pathlib.Path("example-docs").iterdir():
|
||
if file.is_file() and file.suffix == f".{extension}":
|
||
elements = fun(str(file))
|
||
assert all(
|
||
el.metadata.filetype == FILETYPE_TO_MIMETYPE[filetype]
|
||
for el in elements
|
||
if el.metadata.filetype is not None
|
||
)
|
||
break
|
||
|
||
|
||
def test_auto_partition_xml_from_filename(filename="example-docs/factbook.xml"):
|
||
elements = partition(filename=filename, xml_keep_tags=False)
|
||
|
||
assert elements[0].text == "United States"
|
||
assert elements[0].metadata.filename == "factbook.xml"
|
||
|
||
|
||
def test_auto_partition_xml_from_file(filename="example-docs/factbook.xml"):
|
||
with open(filename, "rb") as f:
|
||
elements = partition(file=f, xml_keep_tags=False)
|
||
|
||
assert elements[0].text == "United States"
|
||
|
||
|
||
def test_auto_partition_xml_from_filename_with_tags(filename="example-docs/factbook.xml"):
|
||
elements = partition(filename=filename, xml_keep_tags=True)
|
||
|
||
assert elements[5].text == "<name>United States</name>"
|
||
assert elements[5].metadata.filename == "factbook.xml"
|
||
|
||
|
||
def test_auto_partition_xml_from_file_with_tags(filename="example-docs/factbook.xml"):
|
||
with open(filename, "rb") as f:
|
||
elements = partition(file=f, xml_keep_tags=True)
|
||
|
||
assert elements[5].text == "<name>United States</name>"
|
||
|
||
|
||
EXPECTED_XLSX_FILETYPE = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
|
||
|
||
|
||
def test_auto_partition_xlsx_from_filename(filename="example-docs/stanley-cups.xlsx"):
|
||
elements = partition(filename=filename)
|
||
|
||
assert all(isinstance(element, Table) for element in elements)
|
||
assert len(elements) == 2
|
||
|
||
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TEXT
|
||
assert elements[0].metadata.text_as_html == EXPECTED_TABLE
|
||
assert elements[0].metadata.page_number == 1
|
||
assert elements[0].metadata.filetype == EXPECTED_XLSX_FILETYPE
|
||
|
||
|
||
def test_auto_partition_xlsx_from_file(filename="example-docs/stanley-cups.xlsx"):
|
||
with open(filename, "rb") as f:
|
||
elements = partition(file=f)
|
||
|
||
assert all(isinstance(element, Table) for element in elements)
|
||
assert len(elements) == 2
|
||
|
||
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TEXT
|
||
assert elements[0].metadata.text_as_html == EXPECTED_TABLE
|
||
assert elements[0].metadata.page_number == 1
|
||
assert elements[0].metadata.filetype == EXPECTED_XLSX_FILETYPE
|
||
|
||
|
||
EXPECTED_XLS_TEXT_LEN = 883
|
||
|
||
EXPECTED_XLS_INITIAL_45_CLEAN_TEXT = "MA What C datatypes are 8 bits? (assume i386)"
|
||
|
||
EXPECTED_XLS_TABLE = (
|
||
"""<table border="1" class="dataframe">
|
||
<tbody>
|
||
<tr>
|
||
<td>MA</td>
|
||
<td>What C datatypes are 8 bits? (assume i386)</td>
|
||
<td>int</td>
|
||
<td></td>
|
||
<td>float</td>
|
||
<td></td>
|
||
<td>double</td>
|
||
<td></td>
|
||
<td>char</td>
|
||
</tr>
|
||
<tr>
|
||
<td>TF</td>
|
||
<td>Bagpipes are awesome.</td>
|
||
<td>true</td>
|
||
<td></td>
|
||
<td></td>
|
||
<td></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></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></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></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>
|
||
</tbody>
|
||
</table>"""
|
||
)
|
||
|
||
|
||
def test_auto_partition_xls_from_filename(filename="example-docs/tests-example.xls"):
|
||
elements = partition(filename=filename)
|
||
|
||
assert all(isinstance(element, Table) for element in elements)
|
||
assert len(elements) == 3
|
||
|
||
assert clean_extra_whitespace(elements[0].text)[:45] == EXPECTED_XLS_INITIAL_45_CLEAN_TEXT
|
||
assert len(elements[0].text) == EXPECTED_XLS_TEXT_LEN
|
||
assert elements[0].metadata.text_as_html == EXPECTED_XLS_TABLE
|
||
|
||
|
||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||
def test_auto_partition_csv_from_filename(filename="example-docs/stanley-cups.csv"):
|
||
elements = partition(filename=filename)
|
||
|
||
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TEXT
|
||
assert elements[0].metadata.text_as_html == EXPECTED_TABLE
|
||
assert elements[0].metadata.filetype == "text/csv"
|
||
|
||
|
||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||
def test_auto_partition_csv_from_file(filename="example-docs/stanley-cups.csv"):
|
||
with open(filename, "rb") as f:
|
||
elements = partition(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 == "text/csv"
|
||
|
||
|
||
def test_auto_partition_works_on_empty_filename(filename="example-docs/empty.txt"):
|
||
assert partition(filename=filename) == []
|
||
|
||
|
||
def test_auto_partition_works_on_empty_file(filename="example-docs/empty.txt"):
|
||
with open(filename, "rb") as f:
|
||
assert partition(file=f) == []
|
||
|
||
|
||
def test_auto_partition_rst_from_filename(filename="example-docs/README.rst"):
|
||
elements = partition(filename=filename)
|
||
|
||
assert elements[0] == Title("Example Docs")
|
||
assert elements[0].metadata.filetype == "text/x-rst"
|
||
|
||
|
||
def test_auto_partition_rst_from_file(filename="example-docs/README.rst"):
|
||
with open(filename, "rb") as f:
|
||
elements = partition(file=f, content_type="text/x-rst")
|
||
|
||
assert elements[0] == Title("Example Docs")
|
||
assert elements[0].metadata.filetype == "text/x-rst"
|