unstructured-io--unstructured/test_unstructured/chunking/test_html_output.py

from functools import partial

import pytest

from unstructured.chunking.basic import chunk_elements
from unstructured.chunking.title import chunk_by_title
from unstructured.documents.elements import ElementMetadata, NarrativeText, Text, Title


@pytest.fixture(params=[chunk_elements, partial(chunk_by_title, combine_text_under_n_chars=0)])
def chunking_fn(request):
    return request.param


def test_combining_html_metadata_when_multiple_elements_in_composite_element(chunking_fn):
    metadata_1 = '<h1 class="Title" id="1">Header </h1>'
    metadata_2 = '<time class="CalendarDate" id="2">Date: October 30, 2023 </time>'
    metadata_3 = (
        '<form class="Form" id="3"> '
        '<label class="FormField" for="company-name" id="4">Form field name </label>'
        '<input class="FormFieldValue" id="5" value="Example value" />'
        "</form>"
    )
    combined_metadata = " ".join([metadata_1, metadata_2, metadata_3])

    elements = [
        Title(text="Header", metadata=ElementMetadata(text_as_html=metadata_1)),
        Text(text="Date: October 30, 2023", metadata=ElementMetadata(text_as_html=metadata_2)),
        Text(
            text="Form field name Example value", metadata=ElementMetadata(text_as_html=metadata_3)
        ),
    ]
    chunks = chunking_fn(elements)
    assert len(chunks) == 1
    assert chunks[0].metadata.text_as_html == combined_metadata


def test_combining_html_metadata_with_nested_relationship_between_elements(chunking_fn):
    """
    Ground truth
    <Document>
        <Page>
            <Section>
                <p>First</p>
                <p>Second</p>
            </Section>
        </Page>
    </Document>
    Elements: Document, Page, Section, Paragraph, Paragraph
    Chunk 1: Document, Page, Section, Paragraph

    Chunk 2:
        Paragraph
    """

    metadata_1 = '<div class="Section" id="1" />'
    metadata_2 = '<p class="Paragraph" id="2">First </p>'
    metadata_3 = '<p class="Paragraph" id="3">Second </p>'

    elements = [
        Text(text="", metadata=ElementMetadata(text_as_html=metadata_1)),
        NarrativeText(
            text="First", metadata=ElementMetadata(text_as_html=metadata_2, parent_id="1")
        ),
        NarrativeText(
            text="Second", metadata=ElementMetadata(text_as_html=metadata_3, parent_id="1")
        ),
    ]
    chunks = chunking_fn(elements, max_characters=6)
    assert len(chunks) == 2
    assert chunks[0].text == "First"
    assert chunks[1].text == "Second"

    assert chunks[0].metadata.text_as_html == metadata_1 + " " + metadata_2
    assert chunks[1].metadata.text_as_html == metadata_3


def test_html_metadata_exist_in_both_element_when_text_is_split(chunking_fn):
    """Mimic behaviour of elements with non-html metadata"""
    metadata_1 = '<h1 class="Title" id="1">Header </h1>'
    elements = [
        Title(text="Header", metadata=ElementMetadata(text_as_html=metadata_1)),
    ]
    chunks = chunking_fn(elements, max_characters=3)
    assert len(chunks) == 2

    assert chunks[0].text == "Hea"
    assert chunks[1].text == "der"
    assert chunks[0].metadata.text_as_html == '<h1 class="Title" id="1">Header </h1>'
    assert chunks[1].metadata.text_as_html == '<h1 class="Title" id="1">Header </h1>'