unstructured

1
import pytest
2

3
from unstructured.documents.base import Document, Page
4
from unstructured.documents.elements import Formula, NarrativeText, Title
5

6

7
class MockDocument(Document):
8
    def __init__(self):
9
        super().__init__()
10
        elements = [
11
            Title(text="This is a narrative."),
12
            NarrativeText(text="This is a narrative."),
13
            NarrativeText(text="This is a narrative."),
14
        ]
15
        page = Page(number=0)
16
        page.elements = elements
17
        self._pages = [page]
18

19

20
class MockDocumentWithFormula(Document):
21
    def __init__(self):
22
        super().__init__()
23
        elements = [
24
            Title(text="This is a narrative."),
25
            Formula(text="e=mc2"),
26
        ]
27
        page = Page(number=0)
28
        page.elements = elements
29
        self._pages = [page]
30

31

32
def test_get_narrative():
33
    document = MockDocument()
34
    narrative = document.get_narrative()
35
    for element in narrative:
36
        assert isinstance(element, NarrativeText)
37
    document.print_narrative()
38

39

40
def test_get_formula():
41
    document = MockDocumentWithFormula()
42
    formula = [e for e in document.elements if isinstance(e, Formula)]
43
    assert formula[0].text != ""
44

45

46
@pytest.mark.parametrize("index", [0, 1, 2])
47
def test_split(index):
48
    document = MockDocument()
49
    elements = document.pages[0].elements
50
    split_before_doc = document.before_element(elements[index])
51
    before_elements = split_before_doc.pages[0].elements if split_before_doc.pages else []
52
    split_after_doc = document.after_element(elements[index])
53
    after_elements = split_after_doc.pages[0].elements if split_after_doc.pages else []
54
    expected_before_elements = document.pages[0].elements[:index]
55
    next_index = index + 1
56
    expected_after_elements = document.pages[0].elements[next_index:]
57
    assert all(a.id == b.id for a, b in zip(before_elements, expected_before_elements))
58
    assert all(a.id == b.id for a, b in zip(after_elements, expected_after_elements))
59