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

# pyright: reportPrivateUsage=false

"""Regression and characterization tests for table isolation during chunking (issue #3921).

`Table` (and `TableChunk`, which subclasses `Table`) must never share a pre-chunk with unrelated
text elements, so downstream logic can emit standalone table chunks instead of `CompositeElement`
wrapping mixed content.
"""

from __future__ import annotations

from typing import Callable

import pytest

from unstructured.chunking.base import (
    ChunkingOptions,
    PreChunk,
    PreChunkBuilder,
    PreChunkCombiner,
    PreChunker,
)
from unstructured.chunking.basic import chunk_elements
from unstructured.chunking.title import chunk_by_title
from unstructured.documents.elements import (
    CompositeElement,
    ElementMetadata,
    Table,
    TableChunk,
    Text,
    Title,
)


class DescribeTableIsolationPreChunkBuilder:
    """`PreChunkBuilder` must keep every table-family element in its own pre-chunk."""

    @pytest.mark.parametrize(
        ("preamble", "make_table"),
        [
            (Text("Short preamble."), lambda: Table("H\nC")),
            (Text("Short preamble."), lambda: TableChunk(text="x", metadata=ElementMetadata())),
        ],
    )
    def it_refuses_to_append_a_table_after_any_other_element(
        self, preamble: Text, make_table: Callable[[], Table]
    ):
        builder = PreChunkBuilder(opts=ChunkingOptions(max_characters=500))
        builder.add_element(preamble)

        assert not builder.will_fit(make_table())

    @pytest.mark.parametrize(
        "trailing",
        [
            Text("Follow-up paragraph."),
            Title("Next section"),
        ],
    )
    def it_refuses_to_append_a_non_table_after_a_table(self, trailing: Text | Title):
        builder = PreChunkBuilder(opts=ChunkingOptions(max_characters=500))
        builder.add_element(Table("Heading\nCell"))

        assert not builder.will_fit(trailing)

    def it_allows_only_a_table_when_the_builder_is_empty(self):
        builder = PreChunkBuilder(opts=ChunkingOptions())
        assert builder.will_fit(Table("Heading\nCell text"))

    def it_allows_text_after_flush_even_if_previous_pre_chunk_was_a_table(self):
        opts = ChunkingOptions(max_characters=200)
        builder = PreChunkBuilder(opts=opts)
        builder.add_element(Table("Heading\nCell text"))
        list(builder.flush())  # clears builder state

        assert builder.will_fit(Text("Fresh start after table pre-chunk."))


class DescribeTableIsolationPreChunkStream:
    """End-to-end pre-chunk segmentation from `PreChunker.iter_pre_chunks`."""

    def it_emits_a_table_only_pre_chunk_between_text_blocks(self):
        elements = [
            Title("Section A"),
            Text("Narrative before the table."),
            Table("Col1\nCell A"),
            Text("Narrative after the table."),
        ]
        # -- `new_after_n_chars=0` forces one element per pre-chunk (deterministic layout) --
        opts = ChunkingOptions(max_characters=500, new_after_n_chars=0)

        pre_chunks = list(PreChunker.iter_pre_chunks(elements, opts=opts))

        assert len(pre_chunks) == 4
        assert pre_chunks[0]._elements == [Title("Section A")]
        assert pre_chunks[1]._elements == [Text("Narrative before the table.")]
        assert pre_chunks[2]._elements == [Table("Col1\nCell A")]
        assert pre_chunks[3]._elements == [Text("Narrative after the table.")]

    def it_emits_one_pre_chunk_per_table_when_multiple_tables_are_adjacent(self):
        elements = [
            Table("T1\nA"),
            Table("T2\nB"),
            Text("Closing text."),
        ]
        opts = ChunkingOptions(max_characters=500)

        pre_chunks = list(PreChunker.iter_pre_chunks(elements, opts=opts))

        assert len(pre_chunks) == 3
        assert pre_chunks[0]._elements == [Table("T1\nA")]
        assert pre_chunks[1]._elements == [Table("T2\nB")]
        assert pre_chunks[2]._elements == [Text("Closing text.")]


class DescribeTableIsolationPreChunkCombiner:
    """`PreChunkCombiner` must not stitch table pre-chunks onto text neighbors."""

    def it_keeps_a_table_pre_chunk_separate_when_combining_is_enabled(self):
        opts = ChunkingOptions(max_characters=500, combine_text_under_n_chars=500)
        stream = [
            PreChunk([Text("Hello world.")], overlap_prefix="", opts=opts),
            PreChunk([Table("H\nC")], overlap_prefix="", opts=opts),
            PreChunk([Text("Goodbye world.")], overlap_prefix="", opts=opts),
        ]

        combined = list(PreChunkCombiner(stream, opts=opts).iter_combined_pre_chunks())

        assert len(combined) == 3
        assert combined[0]._elements == [Text("Hello world.")]
        assert combined[1]._elements == [Table("H\nC")]
        assert combined[2]._elements == [Text("Goodbye world.")]


class DescribeTableIsolationOrderingGuarantees:
    """Invariants that should hold for any future refactor of table isolation."""

    def it_preserves_global_element_order_in_pre_chunks(self):
        elements = [
            Text("alpha"),
            Table("T\n1"),
            Text("beta"),
            Table("T\n2"),
            Title("gamma"),
        ]
        opts = ChunkingOptions(max_characters=500, new_after_n_chars=0)
        flat = [e for pc in PreChunker.iter_pre_chunks(elements, opts=opts) for e in pc._elements]

        assert flat == elements

    def it_preserves_global_element_order_in_chunk_elements_output(self):
        elements = [
            Title("Intro"),
            Text("Body before."),
            Table("K\nV"),
            Text("Body after."),
        ]
        chunks = chunk_elements(elements, max_characters=500, new_after_n_chars=0)
        # -- flatten chunk element categories in stream order (table chunks are atomic) --
        categories = [c.category for c in chunks]

        assert categories == [
            "CompositeElement",
            "CompositeElement",
            "Table",
            "CompositeElement",
        ]


class DescribeTableIsolationChunkElements:
    """`chunk_elements` (basic strategy) should emit `Table`/`TableChunk`, not mixed composites."""

    def it_does_not_wrap_a_table_and_surrounding_text_in_one_composite_element(self):
        elements = [
            Title("Report"),
            Text("Short intro."),
            Table("Key\nValue"),
            Text("Short outro."),
        ]
        chunks = chunk_elements(
            elements,
            max_characters=500,
            new_after_n_chars=0,
        )

        assert len(chunks) == 4
        assert isinstance(chunks[0], CompositeElement)
        assert isinstance(chunks[1], CompositeElement)
        assert isinstance(chunks[2], Table)
        assert isinstance(chunks[3], CompositeElement)
        assert "Key" in chunks[2].text or "Value" in chunks[2].text

    def it_yields_distinct_chunks_for_two_tables_in_a_row(self):
        elements = [
            Table("T1\nA"),
            Table("T2\nB"),
        ]
        chunks = chunk_elements(elements, max_characters=500)

        assert len(chunks) == 2
        assert all(isinstance(c, Table) for c in chunks)

    def it_still_isolates_a_table_even_when_the_window_is_very_large(self):
        """Regression: isolation is a semantic rule, not a size heuristic."""
        elements = [
            Text("x"),
            Table("tiny"),
            Text("y"),
        ]
        chunks = chunk_elements(elements, max_characters=50_000, new_after_n_chars=10_000)

        table_chunks = [c for c in chunks if isinstance(c, Table)]
        composite_chunks = [c for c in chunks if isinstance(c, CompositeElement)]

        assert len(table_chunks) == 1
        assert len(composite_chunks) == 2
        assert "tiny" in table_chunks[0].text

    def it_never_produces_a_composite_element_that_lists_a_table_in_orig_elements(
        self,
    ):
        """`CompositeElement` chunks come from `_Chunker`, not `_TableChunker`."""
        elements = [
            Text("preamble"),
            Table("H\nC"),
            Text("post"),
        ]
        chunks = chunk_elements(
            elements,
            max_characters=400,
            new_after_n_chars=0,
            include_orig_elements=True,
        )

        composites = [c for c in chunks if isinstance(c, CompositeElement)]
        for comp in composites:
            orig = comp.metadata.orig_elements or []
            assert not any(isinstance(e, Table) for e in orig)


class DescribeTableIsolationDisabled:
    """When `isolate_table=False`, the pre-#4307 behavior is restored."""

    def it_lets_a_table_share_a_pre_chunk_with_adjacent_text_when_disabled(self):
        opts = ChunkingOptions(max_characters=500, isolate_table=False)
        builder = PreChunkBuilder(opts=opts)
        builder.add_element(Text("Short preamble."))

        assert builder.will_fit(Table("Heading\nCell text"))

    def it_lets_text_follow_a_table_in_the_same_pre_chunk_when_disabled(self):
        opts = ChunkingOptions(max_characters=500, isolate_table=False)
        builder = PreChunkBuilder(opts=opts)
        builder.add_element(Table("Heading\nCell text"))

        assert builder.will_fit(Text("Follow-up paragraph."))

    def it_combines_a_table_pre_chunk_with_text_neighbors_when_disabled(self):
        opts = ChunkingOptions(
            max_characters=500, combine_text_under_n_chars=500, isolate_table=False
        )
        stream = [
            PreChunk([Text("Hello world.")], overlap_prefix="", opts=opts),
            PreChunk([Table("H\nC")], overlap_prefix="", opts=opts),
            PreChunk([Text("Goodbye world.")], overlap_prefix="", opts=opts),
        ]

        combined = list(PreChunkCombiner(stream, opts=opts).iter_combined_pre_chunks())

        assert len(combined) == 1
        assert combined[0]._elements == [
            Text("Hello world."),
            Table("H\nC"),
            Text("Goodbye world."),
        ]

    def it_wraps_a_table_into_a_composite_element_with_neighbors_when_disabled(self):
        """End-to-end: opting out collapses tiny table + text into one CompositeElement."""
        elements = [
            Text("preamble"),
            Table("H\nC"),
            Text("post"),
        ]
        chunks = chunk_elements(
            elements,
            max_characters=500,
            isolate_table=False,
        )

        assert len(chunks) == 1
        assert isinstance(chunks[0], CompositeElement)
        text = chunks[0].text
        assert "preamble" in text
        assert "post" in text


class DescribeTableIsolationOverlapAll:
    """With overlap_all=True, overlap must not cross table / narrative boundaries."""

    def it_does_not_prefix_table_chunk_with_prior_text_overlap(self):
        """Regression: pre-chunk overlap_tail must not become table's overlap_prefix."""
        elements = [Text("Alpha beta gamma delta."), Table("H\nC")]
        chunks = chunk_elements(
            elements,
            max_characters=500,
            new_after_n_chars=0,
            overlap=5,
            overlap_all=True,
        )

        table_chunks = [c for c in chunks if isinstance(c, Table)]
        assert len(table_chunks) == 1
        t = table_chunks[0].text or ""
        assert "Alpha" not in t
        assert "elta" not in t  # tail of "delta" leaked in buggy overlap

    def it_does_not_prefix_text_after_table_with_table_overlap(self):
        elements = [Table("H\nC"), Text("Omega sigma tau upsilon.")]
        chunks = chunk_elements(
            elements,
            max_characters=500,
            new_after_n_chars=0,
            overlap=5,
            overlap_all=True,
        )

        composites = [c for c in chunks if isinstance(c, CompositeElement)]
        assert len(composites) == 1
        assert composites[0].text.startswith("Omega")
        assert "H" not in composites[0].text[:20]

    def it_chunk_by_title_respects_same_overlap_boundaries(self):
        elements = [
            Title("Section"),
            Text("Alpha beta gamma delta."),
            Table("H\nC"),
            Text("Omega sigma tau upsilon."),
        ]
        chunks = chunk_by_title(
            elements,
            max_characters=500,
            new_after_n_chars=0,
            overlap=5,
            overlap_all=True,
            combine_text_under_n_chars=0,
        )

        table_chunks = [c for c in chunks if isinstance(c, Table)]
        assert len(table_chunks) == 1
        assert "Alpha" not in (table_chunks[0].text or "")
        assert "elta" not in (table_chunks[0].text or "")

        omega_composites = [
            c for c in chunks if isinstance(c, CompositeElement) and "Omega" in (c.text or "")
        ]
        assert len(omega_composites) == 1
        assert omega_composites[0].text.startswith("Omega")