unstructured

Форк
0
1291 строка · 46.1 Кб
1
import json
2
import os
3
import pathlib
4
import tempfile
5
import warnings
6
from importlib import import_module
7
from unittest.mock import Mock, patch
8

9
import docx
10
import pytest
11
from PIL import Image
12

13
from test_unstructured.partition.pdf_image.test_pdf import assert_element_extraction
14
from test_unstructured.partition.test_constants import (
15
    EXPECTED_TABLE,
16
    EXPECTED_TABLE_XLSX,
17
    EXPECTED_TEXT,
18
    EXPECTED_TEXT_XLSX,
19
    EXPECTED_TITLE,
20
)
21
from unstructured.chunking.title import chunk_by_title
22
from unstructured.cleaners.core import clean_extra_whitespace
23
from unstructured.documents.elements import (
24
    Address,
25
    ElementMetadata,
26
    ListItem,
27
    NarrativeText,
28
    Table,
29
    TableChunk,
30
    Text,
31
    Title,
32
)
33
from unstructured.file_utils.filetype import FILETYPE_TO_MIMETYPE, FileType
34
from unstructured.partition import auto
35
from unstructured.partition.auto import _get_partition_with_extras, partition
36
from unstructured.partition.common import convert_office_doc
37
from unstructured.partition.utils.constants import PartitionStrategy
38
from unstructured.staging.base import elements_to_json
39

40
DIRECTORY = pathlib.Path(__file__).parent.resolve()
41
EXAMPLE_DOCS_DIRECTORY = os.path.join(DIRECTORY, "..", "..", "example-docs")
42

43
EXPECTED_EMAIL_OUTPUT = [
44
    NarrativeText(text="This is a test email to use for unit tests."),
45
    Title(text="Important points:"),
46
    ListItem(text="Roses are red"),
47
    ListItem(text="Violets are blue"),
48
]
49

50
EML_TEST_FILE = "eml/fake-email.eml"
51

52
is_in_docker = os.path.exists("/.dockerenv")
53

54

55
def test_auto_partition_email_from_filename():
56
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, EML_TEST_FILE)
57
    elements = partition(filename=filename, strategy=PartitionStrategy.HI_RES)
58
    assert len(elements) > 0
59
    assert elements == EXPECTED_EMAIL_OUTPUT
60
    assert elements[0].metadata.filename == os.path.basename(filename)
61
    assert elements[0].metadata.file_directory == os.path.split(filename)[0]
62

63

64
def test_auto_partition_email_from_file():
65
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, EML_TEST_FILE)
66
    with open(filename) as f:
67
        elements = partition(file=f, strategy=PartitionStrategy.HI_RES)
68
    assert len(elements) > 0
69
    assert elements == EXPECTED_EMAIL_OUTPUT
70

71

72
def test_auto_partition_email_from_file_rb():
73
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, EML_TEST_FILE)
74
    with open(filename, "rb") as f:
75
        elements = partition(file=f, strategy=PartitionStrategy.HI_RES)
76
    assert len(elements) > 0
77
    assert elements == EXPECTED_EMAIL_OUTPUT
78

79

80
@pytest.fixture()
81
def mock_docx_document():
82
    document = docx.Document()
83

84
    document.add_paragraph("These are a few of my favorite things:", style="Heading 1")
85
    # NOTE(robinson) - this should get picked up as a list item due to the •
86
    document.add_paragraph("• Parrots", style="Normal")
87
    document.add_paragraph("Hockey", style="List Bullet")
88
    # NOTE(robinson) - this should get picked up as a title
89
    document.add_paragraph("Analysis", style="Normal")
90
    # NOTE(robinson) - this should get dropped because it is empty
91
    document.add_paragraph("", style="Normal")
92
    # NOTE(robinson) - this should get picked up as a narrative text
93
    document.add_paragraph("This is my first thought. This is my second thought.", style="Normal")
94
    document.add_paragraph("This is my third thought.", style="Body Text")
95
    # NOTE(robinson) - this should just be regular text
96
    document.add_paragraph("2023")
97

98
    return document
99

100

101
@pytest.fixture()
102
def expected_docx_elements():
103
    return [
104
        Title("These are a few of my favorite things:"),
105
        ListItem("Parrots"),
106
        ListItem("Hockey"),
107
        Title("Analysis"),
108
        NarrativeText("This is my first thought. This is my second thought."),
109
        NarrativeText("This is my third thought."),
110
        Text("2023"),
111
    ]
112

113

114
def test_auto_partition_docx_with_filename(mock_docx_document, expected_docx_elements, tmpdir):
115
    filename = os.path.join(tmpdir.dirname, "mock_document.docx")
116
    mock_docx_document.save(filename)
117

118
    elements = partition(filename=filename, strategy=PartitionStrategy.HI_RES)
119
    assert elements == expected_docx_elements
120
    assert elements[0].metadata.filename == os.path.basename(filename)
121

122

123
def test_auto_partition_docx_with_file(mock_docx_document, expected_docx_elements, tmpdir):
124
    filename = os.path.join(tmpdir.dirname, "mock_document.docx")
125
    mock_docx_document.save(filename)
126

127
    with open(filename, "rb") as f:
128
        elements = partition(file=f, strategy=PartitionStrategy.HI_RES)
129
    assert elements == expected_docx_elements
130

131

132
@pytest.mark.parametrize(
133
    ("pass_metadata_filename", "content_type"),
134
    [(False, None), (False, "application/msword"), (True, "application/msword"), (True, None)],
135
)
136
def test_auto_partition_doc_with_filename(
137
    mock_docx_document,
138
    expected_docx_elements,
139
    tmpdir,
140
    pass_metadata_filename,
141
    content_type,
142
):
143
    docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
144
    doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
145
    mock_docx_document.save(docx_filename)
146
    convert_office_doc(docx_filename, tmpdir.dirname, "doc")
147
    metadata_filename = doc_filename if pass_metadata_filename else None
148
    elements = partition(
149
        filename=doc_filename,
150
        metadata_filename=metadata_filename,
151
        content_type=content_type,
152
        strategy=PartitionStrategy.HI_RES,
153
    )
154
    assert elements == expected_docx_elements
155
    assert elements[0].metadata.filename == "mock_document.doc"
156
    assert elements[0].metadata.file_directory == tmpdir.dirname
157

158

159
# NOTE(robinson) - the application/x-ole-storage mime type is not specific enough to
160
# determine that the file is an .doc document
161
@pytest.mark.xfail()
162
def test_auto_partition_doc_with_file(mock_docx_document, expected_docx_elements, tmpdir):
163
    docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
164
    doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
165
    mock_docx_document.save(docx_filename)
166
    convert_office_doc(docx_filename, tmpdir.dirname, "doc")
167

168
    with open(doc_filename, "rb") as f:
169
        elements = partition(file=f, strategy=PartitionStrategy.HI_RES)
170
    assert elements == expected_docx_elements
171

172

173
@pytest.mark.parametrize(
174
    ("pass_metadata_filename", "content_type"),
175
    [(False, None), (False, "text/html"), (True, "text/html"), (True, None)],
176
)
177
def test_auto_partition_html_from_filename(pass_metadata_filename, content_type):
178
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "example-10k.html")
179
    metadata_filename = filename if pass_metadata_filename else None
180
    elements = partition(
181
        filename=filename,
182
        metadata_filename=metadata_filename,
183
        content_type=content_type,
184
        strategy=PartitionStrategy.HI_RES,
185
    )
186
    assert len(elements) > 0
187
    assert elements[0].metadata.filename == os.path.basename(filename)
188
    assert elements[0].metadata.file_directory == os.path.split(filename)[0]
189

190

191
@pytest.mark.parametrize(
192
    ("pass_metadata_filename", "content_type"),
193
    [(False, None), (False, "text/html"), (True, "text/html"), (True, None)],
194
)
195
def test_auto_partition_html_from_file(pass_metadata_filename, content_type):
196
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-html.html")
197
    metadata_filename = filename if pass_metadata_filename else None
198
    with open(filename) as f:
199
        elements = partition(
200
            file=f,
201
            metadata_filename=metadata_filename,
202
            content_type=content_type,
203
            strategy=PartitionStrategy.HI_RES,
204
        )
205
    assert len(elements) > 0
206

207

208
def test_auto_partition_html_from_file_rb():
209
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-html.html")
210
    with open(filename, "rb") as f:
211
        elements = partition(file=f, strategy=PartitionStrategy.HI_RES)
212
    assert len(elements) > 0
213

214

215
def test_auto_partition_json_from_filename():
216
    """Test auto-processing an unstructured json output file by filename."""
217
    filename = os.path.join(
218
        EXAMPLE_DOCS_DIRECTORY,
219
        "..",
220
        "test_unstructured_ingest",
221
        "expected-structured-output",
222
        "azure",
223
        "spring-weather.html.json",
224
    )
225
    with open(filename) as json_f:
226
        json_data = json.load(json_f)
227
    json_elems = json.loads(
228
        elements_to_json(partition(filename=filename, strategy=PartitionStrategy.HI_RES))
229
    )
230
    for elem in json_elems:
231
        elem.pop("metadata")
232
    for elem in json_data:
233
        elem.pop("metadata")
234
    assert json_data == json_elems
235

236

237
def test_auto_partition_json_raises_with_unprocessable_json(tmpdir):
238
    # NOTE(robinson) - This is unprocessable because it is not a list of dicts,
239
    # per the Unstructured ISD format
240
    text = '{"hi": "there"}'
241

242
    filename = os.path.join(tmpdir, "unprocessable.json")
243
    with open(filename, "w") as f:
244
        f.write(text)
245

246
    with pytest.raises(ValueError):
247
        partition(filename=filename)
248

249

250
@pytest.mark.xfail(
251
    reason="parsed as text not json, https://github.com/Unstructured-IO/unstructured/issues/492",
252
)
253
def test_auto_partition_json_from_file():
254
    """Test auto-processing an unstructured json output file by file handle."""
255
    filename = os.path.join(
256
        EXAMPLE_DOCS_DIRECTORY,
257
        "..",
258
        "test_unstructured_ingest",
259
        "expected-structured-output",
260
        "azure-blob-storage",
261
        "spring-weather.html.json",
262
    )
263
    with open(filename) as json_f:
264
        json_data = json.load(json_f)
265
    with open(filename, encoding="utf-8") as partition_f:
266
        json_elems = json.loads(
267
            elements_to_json(partition(file=partition_f, strategy=PartitionStrategy.HI_RES))
268
        )
269
    for elem in json_elems:
270
        # coordinates are always in the element data structures, even if None
271
        elem.pop("coordinates")
272
        elem.pop("coordinate_system")
273
    assert json_data == json_elems
274

275

276
EXPECTED_TEXT_OUTPUT = [
277
    NarrativeText(text="This is a test document to use for unit tests."),
278
    Address(text="Doylestown, PA 18901"),
279
    Title(text="Important points:"),
280
    ListItem(text="Hamburgers are delicious"),
281
    ListItem(text="Dogs are the best"),
282
    ListItem(text="I love fuzzy blankets"),
283
]
284

285

286
def test_auto_partition_text_from_filename():
287
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-text.txt")
288
    elements = partition(filename=filename, strategy=PartitionStrategy.HI_RES)
289
    assert len(elements) > 0
290
    assert elements == EXPECTED_TEXT_OUTPUT
291
    assert elements[0].metadata.filename == os.path.basename(filename)
292
    assert elements[0].metadata.file_directory == os.path.split(filename)[0]
293

294

295
def test_auto_partition_text_from_file():
296
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-text.txt")
297
    with open(filename) as f:
298
        elements = partition(file=f, strategy=PartitionStrategy.HI_RES)
299
    assert len(elements) > 0
300
    assert elements == EXPECTED_TEXT_OUTPUT
301

302

303
@pytest.mark.parametrize(
304
    ("pass_metadata_filename", "content_type"),
305
    [(False, None), (False, "application/pdf"), (True, "application/pdf"), (True, None)],
306
)
307
def test_auto_partition_pdf_from_filename(pass_metadata_filename, content_type, request):
308
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.pdf")
309
    metadata_filename = filename if pass_metadata_filename else None
310

311
    elements = partition(
312
        filename=filename,
313
        metadata_filename=metadata_filename,
314
        content_type=content_type,
315
        strategy=PartitionStrategy.HI_RES,
316
    )
317

318
    idx = 3
319
    assert isinstance(elements[idx], Title)
320
    assert elements[idx].text.startswith("LayoutParser")
321

322
    assert elements[idx].metadata.filename == os.path.basename(filename)
323
    assert elements[idx].metadata.file_directory == os.path.split(filename)[0]
324

325
    # NOTE(alan): Xfail since new model skips the word Zejiang
326
    request.applymarker(pytest.mark.xfail)
327

328
    idx += 1
329
    assert isinstance(elements[idx], NarrativeText)
330
    assert elements[idx].text.startswith("Zejiang Shen")
331

332

333
def test_auto_partition_pdf_uses_table_extraction():
334
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.pdf")
335
    with patch(
336
        "unstructured.partition.pdf_image.ocr.process_file_with_ocr",
337
    ) as mock_process_file_with_model:
338
        partition(filename, pdf_infer_table_structure=True, strategy=PartitionStrategy.HI_RES)
339
        assert mock_process_file_with_model.call_args[1]["infer_table_structure"]
340

341

342
def test_auto_partition_pdf_with_fast_strategy(monkeypatch):
343
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.pdf")
344

345
    mock_return = [NarrativeText("Hello there!")]
346
    with patch.object(auto, "partition_pdf", return_value=mock_return) as mock_partition:
347
        mock_partition_with_extras_map = {"pdf": mock_partition}
348
        monkeypatch.setattr(auto, "PARTITION_WITH_EXTRAS_MAP", mock_partition_with_extras_map)
349
        partition(filename=filename, strategy=PartitionStrategy.FAST)
350

351
    mock_partition.assert_called_once_with(
352
        filename=filename,
353
        file=None,
354
        url=None,
355
        strategy=PartitionStrategy.FAST,
356
        languages=None,
357
        metadata_filename=None,
358
        include_page_breaks=False,
359
        infer_table_structure=False,
360
        extract_images_in_pdf=False,
361
        extract_image_block_types=None,
362
        extract_image_block_output_dir=None,
363
        extract_image_block_to_payload=False,
364
        hi_res_model_name=None,
365
    )
366

367

368
@pytest.mark.parametrize(
369
    ("pass_metadata_filename", "content_type"),
370
    [(False, None), (False, "application/pdf"), (True, "application/pdf"), (True, None)],
371
)
372
def test_auto_partition_pdf_from_file(pass_metadata_filename, content_type, request):
373
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.pdf")
374
    metadata_filename = filename if pass_metadata_filename else None
375

376
    with open(filename, "rb") as f:
377
        elements = partition(
378
            file=f,
379
            metadata_filename=metadata_filename,
380
            content_type=content_type,
381
            strategy=PartitionStrategy.HI_RES,
382
        )
383

384
    idx = 3
385
    assert isinstance(elements[idx], Title)
386
    assert elements[idx].text.startswith("LayoutParser")
387

388
    # NOTE(alan): Xfail since new model misses the first word Zejiang
389
    request.applymarker(pytest.mark.xfail)
390

391
    idx += 1
392
    assert isinstance(elements[idx], NarrativeText)
393
    assert elements[idx].text.startswith("Zejiang Shen")
394

395

396
def test_auto_partition_formats_languages_for_tesseract():
397
    filename = "example-docs/chi_sim_image.jpeg"
398
    with patch(
399
        "unstructured.partition.pdf_image.ocr.process_file_with_ocr",
400
    ) as mock_process_file_with_ocr:
401
        partition(filename, strategy=PartitionStrategy.HI_RES, languages=["zh"])
402
        _, kwargs = mock_process_file_with_ocr.call_args_list[0]
403
        assert "ocr_languages" in kwargs
404
        assert kwargs["ocr_languages"] == "chi_sim+chi_sim_vert+chi_tra+chi_tra_vert"
405

406

407
def test_auto_partition_element_metadata_user_provided_languages():
408
    filename = "example-docs/chevron-page.pdf"
409
    elements = partition(filename=filename, strategy=PartitionStrategy.OCR_ONLY, languages=["eng"])
410
    assert elements[0].metadata.languages == ["eng"]
411

412

413
@pytest.mark.parametrize(
414
    ("languages", "ocr_languages"),
415
    [(["auto"], ""), (["eng"], "")],
416
)
417
def test_auto_partition_ignores_empty_string_for_ocr_languages(languages, ocr_languages):
418
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "book-war-and-peace-1p.txt")
419
    elements = partition(
420
        filename=filename,
421
        strategy=PartitionStrategy.OCR_ONLY,
422
        ocr_languages=ocr_languages,
423
        languages=languages,
424
    )
425
    assert elements[0].metadata.languages == ["eng"]
426

427

428
def test_auto_partition_warns_with_ocr_languages(caplog):
429
    filename = "example-docs/chevron-page.pdf"
430
    partition(filename=filename, strategy=PartitionStrategy.HI_RES, ocr_languages="eng")
431
    assert "The ocr_languages kwarg will be deprecated" in caplog.text
432

433

434
def test_partition_pdf_doesnt_raise_warning():
435
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.pdf")
436
    # NOTE(robinson): This is the recommended way to check that no warning is emitted,
437
    # per the pytest docs.
438
    # ref: https://docs.pytest.org/en/7.0.x/how-to/capture-warnings.html
439
    #      #additional-use-cases-of-warnings-in-tests
440
    with warnings.catch_warnings():
441
        warnings.simplefilter("error")
442
        partition(filename=filename, strategy=PartitionStrategy.HI_RES)
443

444

445
@pytest.mark.parametrize(
446
    ("pass_metadata_filename", "content_type"),
447
    [(False, None), (False, "image/jpeg"), (True, "image/jpeg"), (True, None)],
448
)
449
def test_auto_partition_image(pass_metadata_filename, content_type):
450
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.jpg")
451
    metadata_filename = filename if pass_metadata_filename else None
452
    elements = partition(
453
        filename=filename,
454
        metadata_filename=metadata_filename,
455
        content_type=content_type,
456
        strategy=PartitionStrategy.AUTO,
457
    )
458

459
    # should be same result as test_partition_image_default_strategy_hi_res() in test_image.py
460
    title = "LayoutParser: A Unified Toolkit for Deep Learning Based Document Image Analysis"
461
    idx = 2
462
    assert elements[idx].text == title
463
    assert elements[idx].metadata.coordinates is not None
464

465

466
@pytest.mark.parametrize("extract_image_block_to_payload", [False, True])
467
def test_auto_partition_image_element_extraction(
468
    extract_image_block_to_payload,
469
    filename=os.path.join(EXAMPLE_DOCS_DIRECTORY, "embedded-images-tables.jpg"),
470
):
471
    extract_image_block_types = ["Image", "Table"]
472

473
    with tempfile.TemporaryDirectory() as tmpdir:
474
        elements = partition(
475
            filename=filename,
476
            extract_image_block_types=extract_image_block_types,
477
            extract_image_block_to_payload=extract_image_block_to_payload,
478
            extract_image_block_output_dir=tmpdir,
479
        )
480

481
        assert_element_extraction(
482
            elements, extract_image_block_types, extract_image_block_to_payload, tmpdir
483
        )
484

485

486
@pytest.mark.parametrize(
487
    ("pass_metadata_filename", "content_type"),
488
    [(False, None), (False, "image/jpeg"), (True, "image/jpeg"), (True, None)],
489
)
490
def test_auto_partition_jpg(pass_metadata_filename, content_type):
491
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.jpg")
492
    metadata_filename = filename if pass_metadata_filename else None
493
    elements = partition(
494
        filename=filename,
495
        metadata_filename=metadata_filename,
496
        content_type=content_type,
497
        strategy=PartitionStrategy.AUTO,
498
    )
499
    assert len(elements) > 0
500

501

502
@pytest.mark.parametrize(
503
    ("pass_metadata_filename", "content_type"),
504
    [(False, None), (False, "image/jpeg"), (True, "image/jpeg"), (True, None)],
505
)
506
def test_auto_partition_jpg_from_file(pass_metadata_filename, content_type):
507
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.jpg")
508
    metadata_filename = filename if pass_metadata_filename else None
509
    with open(filename, "rb") as f:
510
        elements = partition(
511
            file=f,
512
            metadata_filename=metadata_filename,
513
            content_type=content_type,
514
            strategy=PartitionStrategy.AUTO,
515
        )
516
    assert len(elements) > 0
517

518

519
def test_auto_partition_raises_with_bad_type(monkeypatch):
520
    monkeypatch.setattr(auto, "detect_filetype", lambda *args, **kwargs: None)
521
    with pytest.raises(ValueError):
522
        partition(filename="made-up.fake", strategy=PartitionStrategy.HI_RES)
523

524

525
EXPECTED_PPTX_OUTPUT = [
526
    Title(text="Adding a Bullet Slide"),
527
    ListItem(text="Find the bullet slide layout"),
528
    ListItem(text="Use _TextFrame.text for first bullet"),
529
    ListItem(text="Use _TextFrame.add_paragraph() for subsequent bullets"),
530
    NarrativeText(text="Here is a lot of text!"),
531
    NarrativeText(text="Here is some text in a text box!"),
532
]
533

534

535
def test_auto_partition_pptx_from_filename():
536
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-power-point.pptx")
537
    elements = partition(filename=filename, strategy=PartitionStrategy.HI_RES)
538
    assert elements == EXPECTED_PPTX_OUTPUT
539
    assert elements[0].metadata.filename == os.path.basename(filename)
540
    assert elements[0].metadata.file_directory == os.path.split(filename)[0]
541

542

543
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
544
def test_auto_partition_ppt_from_filename():
545
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-power-point.ppt")
546
    elements = partition(filename=filename, strategy=PartitionStrategy.HI_RES)
547
    assert elements == EXPECTED_PPTX_OUTPUT
548
    assert elements[0].metadata.filename == os.path.basename(filename)
549
    assert elements[0].metadata.file_directory == os.path.split(filename)[0]
550

551

552
def test_auto_with_page_breaks():
553
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper-fast.pdf")
554
    elements = partition(
555
        filename=filename, include_page_breaks=True, strategy=PartitionStrategy.HI_RES
556
    )
557
    assert "PageBreak" in [elem.category for elem in elements]
558

559

560
def test_auto_partition_epub_from_filename():
561
    filename = os.path.join(DIRECTORY, "..", "..", "example-docs", "winter-sports.epub")
562
    elements = partition(filename=filename, strategy=PartitionStrategy.HI_RES)
563
    assert len(elements) > 0
564
    assert elements[0].text.startswith("The Project Gutenberg eBook of Winter Sports")
565

566

567
def test_auto_partition_epub_from_file():
568
    filename = os.path.join(DIRECTORY, "..", "..", "example-docs", "winter-sports.epub")
569
    with open(filename, "rb") as f:
570
        elements = partition(file=f, strategy=PartitionStrategy.HI_RES)
571
    assert len(elements) > 0
572
    assert elements[0].text.startswith("The Project Gutenberg eBook of Winter Sports")
573

574

575
EXPECTED_MSG_OUTPUT = [
576
    NarrativeText(text="This is a test email to use for unit tests."),
577
    Title(text="Important points:"),
578
    ListItem(text="Roses are red"),
579
    ListItem(text="Violets are blue"),
580
]
581

582

583
def test_auto_partition_msg_from_filename():
584
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-email.msg")
585
    elements = partition(filename=filename, strategy=PartitionStrategy.HI_RES)
586
    assert elements == EXPECTED_MSG_OUTPUT
587

588

589
def test_auto_partition_rtf_from_filename():
590
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-doc.rtf")
591
    elements = partition(filename=filename, strategy=PartitionStrategy.HI_RES)
592
    assert elements[0] == Title("My First Heading")
593

594

595
def test_auto_partition_from_url():
596
    url = "https://raw.githubusercontent.com/Unstructured-IO/unstructured/main/LICENSE.md"
597
    elements = partition(url=url, content_type="text/plain", strategy=PartitionStrategy.HI_RES)
598
    assert elements[0] == Title("Apache License")
599
    assert elements[0].metadata.url == url
600

601

602
def test_partition_md_works_with_embedded_html():
603
    url = "https://raw.githubusercontent.com/Unstructured-IO/unstructured/main/README.md"
604
    elements = partition(url=url, content_type="text/markdown", strategy=PartitionStrategy.HI_RES)
605
    elements[0].text
606
    unstructured_found = False
607
    for element in elements:
608
        if "unstructured" in elements[0].text:
609
            unstructured_found = True
610
            break
611
    assert unstructured_found is True
612

613

614
def test_auto_partition_warns_if_header_set_and_not_url(caplog):
615
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, EML_TEST_FILE)
616
    partition(
617
        filename=filename, headers={"Accept": "application/pdf"}, strategy=PartitionStrategy.HI_RES
618
    )
619
    assert caplog.records[0].levelname == "WARNING"
620

621

622
def test_auto_partition_works_with_unstructured_jsons():
623
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "spring-weather.html.json")
624
    elements = partition(filename=filename, strategy=PartitionStrategy.HI_RES)
625
    assert elements[0].text == "News Around NOAA"
626

627

628
def test_auto_partition_works_with_unstructured_jsons_from_file():
629
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "spring-weather.html.json")
630
    with open(filename, "rb") as f:
631
        elements = partition(file=f, strategy=PartitionStrategy.HI_RES)
632
    assert elements[0].text == "News Around NOAA"
633

634

635
def test_auto_partition_odt_from_filename():
636
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake.odt")
637
    elements = partition(filename=filename, strategy=PartitionStrategy.HI_RES)
638
    assert elements[0] == Title("Lorem ipsum dolor sit amet.")
639

640

641
def test_auto_partition_odt_from_file():
642
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake.odt")
643
    with open(filename, "rb") as f:
644
        elements = partition(file=f, strategy=PartitionStrategy.HI_RES)
645

646
    assert elements[0] == Title("Lorem ipsum dolor sit amet.")
647

648

649
@pytest.mark.parametrize(
650
    ("content_type", "routing_func", "expected"),
651
    [
652
        ("text/csv", "csv", "text/csv"),
653
        ("text/html", "html", "text/html"),
654
        ("jdsfjdfsjkds", "pdf", None),
655
    ],
656
)
657
def test_auto_adds_filetype_to_metadata(content_type, routing_func, expected, monkeypatch):
658
    with patch(
659
        f"unstructured.partition.auto.partition_{routing_func}",
660
        lambda *args, **kwargs: [Text("text 1"), Text("text 2")],
661
    ) as mock_partition:
662
        mock_partition_with_extras_map = {routing_func: mock_partition}
663
        monkeypatch.setattr(auto, "PARTITION_WITH_EXTRAS_MAP", mock_partition_with_extras_map)
664
        elements = partition("example-docs/layout-parser-paper-fast.pdf", content_type=content_type)
665
    assert len(elements) == 2
666
    assert all(el.metadata.filetype == expected for el in elements)
667

668

669
@pytest.mark.parametrize(
670
    ("content_type", "expected"),
671
    [
672
        ("application/pdf", FILETYPE_TO_MIMETYPE[FileType.PDF]),
673
        (None, FILETYPE_TO_MIMETYPE[FileType.PDF]),
674
    ],
675
)
676
def test_auto_filetype_overrides_file_specific(content_type, expected, monkeypatch):
677
    pdf_metadata = ElementMetadata(filetype="imapdf")
678
    with patch(
679
        "unstructured.partition.auto.partition_pdf",
680
        lambda *args, **kwargs: [
681
            Text("text 1", metadata=pdf_metadata),
682
            Text("text 2", metadata=pdf_metadata),
683
        ],
684
    ) as mock_partition:
685
        mock_partition_with_extras_map = {"pdf": mock_partition}
686
        monkeypatch.setattr(auto, "PARTITION_WITH_EXTRAS_MAP", mock_partition_with_extras_map)
687
        elements = partition("example-docs/layout-parser-paper-fast.pdf", content_type=content_type)
688
    assert len(elements) == 2
689
    assert all(el.metadata.filetype == expected for el in elements)
690

691

692
@pytest.mark.parametrize("extract_image_block_to_payload", [False, True])
693
def test_auto_partition_pdf_element_extraction(
694
    extract_image_block_to_payload,
695
    filename=os.path.join(EXAMPLE_DOCS_DIRECTORY, "embedded-images-tables.pdf"),
696
):
697
    extract_image_block_types = ["Image", "Table"]
698

699
    with tempfile.TemporaryDirectory() as tmpdir:
700
        elements = partition(
701
            filename=filename,
702
            extract_image_block_types=extract_image_block_types,
703
            extract_image_block_to_payload=extract_image_block_to_payload,
704
            extract_image_block_output_dir=tmpdir,
705
        )
706

707
        assert_element_extraction(
708
            elements, extract_image_block_types, extract_image_block_to_payload, tmpdir
709
        )
710

711

712
supported_filetypes = [
713
    _
714
    for _ in FileType
715
    if _
716
    not in (
717
        FileType.UNK,
718
        FileType.ZIP,
719
        FileType.XLS,
720
    )
721
]
722

723

724
FILETYPE_TO_MODULE = {
725
    FileType.JPG: "image",
726
    FileType.PNG: "image",
727
    FileType.HEIC: "image",
728
    FileType.TXT: "text",
729
    FileType.EML: "email",
730
}
731

732

733
@pytest.mark.parametrize("filetype", supported_filetypes)
734
def test_file_specific_produces_correct_filetype(filetype: FileType):
735
    if filetype in auto.IMAGE_FILETYPES or filetype in (FileType.WAV, FileType.EMPTY):
736
        pytest.skip()
737
    extension = filetype.name.lower()
738
    filetype_module = FILETYPE_TO_MODULE.get(filetype, extension)
739
    fun_name = "partition_" + filetype_module
740
    module = import_module(f"unstructured.partition.{filetype_module}")  # noqa
741
    fun = eval(f"module.{fun_name}")
742
    for file in pathlib.Path("example-docs").iterdir():
743
        if file.is_file() and file.suffix == f".{extension}":
744
            elements = fun(str(file))
745
            assert all(
746
                el.metadata.filetype == FILETYPE_TO_MIMETYPE[filetype]
747
                for el in elements
748
                if el.metadata.filetype is not None
749
            )
750
            break
751

752

753
def test_auto_partition_xml_from_filename(filename="example-docs/factbook.xml"):
754
    elements = partition(filename=filename, xml_keep_tags=False, metadata_filename=filename)
755

756
    assert elements[0].text == "United States"
757
    assert elements[0].metadata.filename == "factbook.xml"
758

759

760
def test_auto_partition_xml_from_file(filename="example-docs/factbook.xml"):
761
    with open(filename, "rb") as f:
762
        elements = partition(file=f, xml_keep_tags=False)
763

764
    assert elements[0].text == "United States"
765

766

767
def test_auto_partition_xml_from_filename_with_tags(filename="example-docs/factbook.xml"):
768
    elements = partition(filename=filename, xml_keep_tags=True)
769

770
    assert "<leader>Joe Biden</leader>" in elements[0].text
771
    assert elements[0].metadata.filename == "factbook.xml"
772

773

774
def test_auto_partition_xml_from_file_with_tags(filename="example-docs/factbook.xml"):
775
    with open(filename, "rb") as f:
776
        elements = partition(file=f, xml_keep_tags=True)
777

778
    assert "<leader>Joe Biden</leader>" in elements[0].text
779

780

781
EXPECTED_XLSX_FILETYPE = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
782

783

784
def test_auto_partition_xlsx_from_filename(filename="example-docs/stanley-cups.xlsx"):
785
    elements = partition(filename=filename, include_header=False, skip_infer_table_types=[])
786

787
    assert sum(isinstance(element, Table) for element in elements) == 2
788
    assert sum(isinstance(element, Title) for element in elements) == 2
789
    assert len(elements) == 4
790

791
    assert clean_extra_whitespace(elements[0].text) == EXPECTED_TITLE
792
    assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
793
    assert elements[1].metadata.text_as_html == EXPECTED_TABLE_XLSX
794
    assert elements[1].metadata.page_number == 1
795
    assert elements[1].metadata.filetype == EXPECTED_XLSX_FILETYPE
796

797

798
@pytest.mark.parametrize(
799
    ("skip_infer_table_types", "filename", "has_text_as_html_field"),
800
    [
801
        (["xlsx"], "stanley-cups.xlsx", False),
802
        ([], "stanley-cups.xlsx", True),
803
        (["odt"], "fake.odt", False),
804
        ([], "fake.odt", True),
805
    ],
806
)
807
def test_auto_partition_respects_skip_infer_table_types(
808
    skip_infer_table_types,
809
    filename,
810
    has_text_as_html_field,
811
):
812
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, filename)
813
    with open(filename, "rb") as f:
814
        table_elements = [
815
            e
816
            for e in partition(file=f, skip_infer_table_types=skip_infer_table_types)
817
            if isinstance(e, Table)
818
        ]
819
        for table_element in table_elements:
820
            table_element_has_text_as_html_field = (
821
                hasattr(table_element.metadata, "text_as_html")
822
                and table_element.metadata.text_as_html is not None
823
            )
824
        assert table_element_has_text_as_html_field == has_text_as_html_field
825

826

827
def test_auto_partition_xlsx_from_file(filename="example-docs/stanley-cups.xlsx"):
828
    with open(filename, "rb") as f:
829
        elements = partition(file=f, include_header=False, skip_infer_table_types=[])
830

831
    assert sum(isinstance(element, Table) for element in elements) == 2
832
    assert sum(isinstance(element, Title) for element in elements) == 2
833
    assert len(elements) == 4
834

835
    assert clean_extra_whitespace(elements[0].text) == EXPECTED_TITLE
836
    assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
837
    assert elements[1].metadata.text_as_html == EXPECTED_TABLE_XLSX
838
    assert elements[1].metadata.page_number == 1
839
    assert elements[1].metadata.filetype == EXPECTED_XLSX_FILETYPE
840

841

842
EXPECTED_XLS_TEXT_LEN = 550
843

844

845
EXPECTED_XLS_INITIAL_45_CLEAN_TEXT = "MC What is 2+2? 4 correct 3 incorrect MA What"
846

847
EXPECTED_XLS_TABLE = (
848
    """<table border="1" class="dataframe">
849
  <tbody>
850
    <tr>
851
      <td>MC</td>
852
      <td>What is 2+2?</td>
853
      <td>4</td>
854
      <td>correct</td>
855
      <td>3</td>
856
      <td>incorrect</td>
857
      <td></td>
858
      <td></td>
859
      <td></td>
860
    </tr>
861
    <tr>
862
      <td>MA</td>
863
      <td>What C datatypes are 8 bits? (assume i386)</td>
864
      <td>int</td>
865
      <td></td>
866
      <td>float</td>
867
      <td></td>
868
      <td>double</td>
869
      <td></td>
870
      <td>char</td>
871
    </tr>
872
    <tr>
873
      <td>TF</td>
874
      <td>Bagpipes are awesome.</td>
875
      <td>true</td>
876
      <td></td>
877
      <td></td>
878
      <td></td>
879
      <td></td>
880
      <td></td>
881
      <td></td>
882
    </tr>
883
    <tr>
884
      <td>ESS</td>
885
      <td>How have the original Henry Hornbostel buildings """
886
    """influenced campus architecture and design in the last 30 years?</td>
887
      <td></td>
888
      <td></td>
889
      <td></td>
890
      <td></td>
891
      <td></td>
892
      <td></td>
893
      <td></td>
894
    </tr>
895
    <tr>
896
      <td>ORD</td>
897
      <td>Rank the following in their order of operation.</td>
898
      <td>Parentheses</td>
899
      <td>Exponents</td>
900
      <td>Division</td>
901
      <td>Addition</td>
902
      <td></td>
903
      <td></td>
904
      <td></td>
905
    </tr>
906
    <tr>
907
      <td>FIB</td>
908
      <td>The student activities fee is</td>
909
      <td>95</td>
910
      <td>dollars for students enrolled in</td>
911
      <td>19</td>
912
      <td>units or more,</td>
913
      <td></td>
914
      <td></td>
915
      <td></td>
916
    </tr>
917
    <tr>
918
      <td>MAT</td>
919
      <td>Match the lower-case greek letter with its capital form.</td>
920
      <td>λ</td>
921
      <td>Λ</td>
922
      <td>α</td>
923
      <td>γ</td>
924
      <td>Γ</td>
925
      <td>φ</td>
926
      <td>Φ</td>
927
    </tr>
928
  </tbody>
929
</table>"""
930
)
931

932

933
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
934
def test_auto_partition_xls_from_filename(filename="example-docs/tests-example.xls"):
935
    elements = partition(filename=filename, include_header=False, skip_infer_table_types=[])
936

937
    assert sum(isinstance(element, Table) for element in elements) == 2
938
    assert len(elements) == 14
939

940
    assert clean_extra_whitespace(elements[0].text)[:45] == EXPECTED_XLS_INITIAL_45_CLEAN_TEXT
941
    # NOTE(crag): if the beautifulsoup4 package is installed, some (but not all) additional
942
    # whitespace is removed, so the expected text length is less than is the case
943
    # when beautifulsoup4 is *not* installed. E.g.
944
    # "\n\n\nMA\nWhat C datatypes are 8 bits" vs.
945
    # '\n  \n    \n      MA\n      What C datatypes are 8 bits?... "
946
    assert len(elements[0].text) == EXPECTED_XLS_TEXT_LEN
947
    assert elements[0].metadata.text_as_html == EXPECTED_XLS_TABLE
948

949

950
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
951
def test_auto_partition_csv_from_filename(filename="example-docs/stanley-cups.csv"):
952
    elements = partition(filename=filename)
953

954
    assert clean_extra_whitespace(elements[0].text) == EXPECTED_TEXT
955
    assert elements[0].metadata.text_as_html == EXPECTED_TABLE
956
    assert elements[0].metadata.filetype == "text/csv"
957

958

959
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
960
def test_auto_partition_tsv_from_filename(filename="example-docs/stanley-cups.tsv"):
961
    elements = partition(filename=filename)
962

963
    assert clean_extra_whitespace(elements[0].text) == EXPECTED_TEXT
964
    assert elements[0].metadata.text_as_html == EXPECTED_TABLE
965
    assert elements[0].metadata.filetype == "text/tsv"
966

967

968
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
969
def test_auto_partition_csv_from_file(filename="example-docs/stanley-cups.csv"):
970
    with open(filename, "rb") as f:
971
        elements = partition(file=f)
972

973
    assert clean_extra_whitespace(elements[0].text) == EXPECTED_TEXT
974
    assert isinstance(elements[0], Table)
975
    assert elements[0].metadata.text_as_html == EXPECTED_TABLE
976
    assert elements[0].metadata.filetype == "text/csv"
977

978

979
def test_auto_partition_html_pre_from_file(filename="example-docs/fake-html-pre.htm"):
980
    elements = partition(filename=filename)
981

982
    assert len(elements) > 0
983
    assert "PageBreak" not in [elem.category for elem in elements]
984
    assert clean_extra_whitespace(elements[0].text).startswith("[107th Congress Public Law 56]")
985
    assert isinstance(elements[0], NarrativeText)
986
    assert elements[0].metadata.filetype == "text/html"
987
    assert elements[0].metadata.filename == "fake-html-pre.htm"
988

989

990
def test_auto_partition_works_on_empty_filename(filename="example-docs/empty.txt"):
991
    assert partition(filename=filename) == []
992

993

994
def test_auto_partition_works_on_empty_file(filename="example-docs/empty.txt"):
995
    with open(filename, "rb") as f:
996
        assert partition(file=f) == []
997

998

999
def test_auto_partition_org_from_filename(filename="example-docs/README.org"):
1000
    elements = partition(filename=filename)
1001

1002
    assert elements[0] == Title("Example Docs")
1003
    assert elements[0].metadata.filetype == "text/org"
1004

1005

1006
def test_auto_partition_org_from_file(filename="example-docs/README.org"):
1007
    with open(filename, "rb") as f:
1008
        elements = partition(file=f, content_type="text/org")
1009

1010
    assert elements[0] == Title("Example Docs")
1011
    assert elements[0].metadata.filetype == "text/org"
1012

1013

1014
def test_auto_partition_rst_from_filename(filename="example-docs/README.rst"):
1015
    elements = partition(filename=filename)
1016

1017
    assert elements[0] == Title("Example Docs")
1018
    assert elements[0].metadata.filetype == "text/x-rst"
1019

1020

1021
def test_auto_partition_rst_from_file(filename="example-docs/README.rst"):
1022
    with open(filename, "rb") as f:
1023
        elements = partition(file=f, content_type="text/x-rst")
1024

1025
    assert elements[0] == Title("Example Docs")
1026
    assert elements[0].metadata.filetype == "text/x-rst"
1027

1028

1029
def test_auto_partition_metadata_filename():
1030
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-text.txt")
1031
    with open(filename) as f:
1032
        elements = partition(file=f, metadata_filename=filename)
1033
    assert elements[0].metadata.filename == os.path.split(filename)[-1]
1034

1035

1036
def test_auto_partition_warns_about_file_filename_deprecation(caplog):
1037
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-text.txt")
1038
    with open(filename) as f:
1039
        elements = partition(file=f, file_filename=filename)
1040
    assert elements[0].metadata.filename == os.path.split(filename)[-1]
1041
    assert "WARNING" in caplog.text
1042
    assert "The file_filename kwarg will be deprecated" in caplog.text
1043

1044

1045
def test_auto_partition_raises_with_file_and_metadata_filename():
1046
    filename = os.path.join(EXAMPLE_DOCS_DIRECTORY, "fake-text.txt")
1047
    with open(filename) as f, pytest.raises(ValueError):
1048
        partition(file=f, file_filename=filename, metadata_filename=filename)
1049

1050

1051
def test_get_partition_with_extras_prompts_for_install_if_missing():
1052
    partition_with_extras_map = {}
1053
    with pytest.raises(ImportError) as exception_info:
1054
        _get_partition_with_extras("pdf", partition_with_extras_map)
1055

1056
    msg = str(exception_info.value)
1057
    assert 'Install the pdf dependencies with pip install "unstructured[pdf]"' in msg
1058

1059

1060
def test_add_chunking_strategy_on_partition_auto():
1061
    filename = "example-docs/example-10k-1p.html"
1062
    elements = partition(filename)
1063
    chunk_elements = partition(filename, chunking_strategy="by_title")
1064
    chunks = chunk_by_title(elements)
1065
    assert chunk_elements != elements
1066
    assert chunk_elements == chunks
1067

1068

1069
def test_add_chunking_strategy_title_on_partition_auto_respects_multipage():
1070
    filename = "example-docs/example-10k-1p.html"
1071
    partitioned_elements_multipage_false_combine_chars_0 = partition(
1072
        filename,
1073
        chunking_strategy="by_title",
1074
        multipage_sections=False,
1075
        combine_text_under_n_chars=0,
1076
        new_after_n_chars=300,
1077
        max_characters=400,
1078
    )
1079
    partitioned_elements_multipage_true_combine_chars_0 = partition(
1080
        filename,
1081
        chunking_strategy="by_title",
1082
        multipage_sections=True,
1083
        combine_text_under_n_chars=0,
1084
        new_after_n_chars=300,
1085
        max_characters=400,
1086
    )
1087
    elements = partition(filename)
1088
    cleaned_elements_multipage_false_combine_chars_0 = chunk_by_title(
1089
        elements,
1090
        multipage_sections=False,
1091
        combine_text_under_n_chars=0,
1092
        new_after_n_chars=300,
1093
        max_characters=400,
1094
    )
1095
    cleaned_elements_multipage_true_combine_chars_0 = chunk_by_title(
1096
        elements,
1097
        multipage_sections=True,
1098
        combine_text_under_n_chars=0,
1099
        new_after_n_chars=300,
1100
        max_characters=400,
1101
    )
1102
    assert (
1103
        partitioned_elements_multipage_false_combine_chars_0
1104
        == cleaned_elements_multipage_false_combine_chars_0
1105
    )
1106
    assert (
1107
        partitioned_elements_multipage_true_combine_chars_0
1108
        == cleaned_elements_multipage_true_combine_chars_0
1109
    )
1110
    assert len(partitioned_elements_multipage_true_combine_chars_0) != len(
1111
        partitioned_elements_multipage_false_combine_chars_0,
1112
    )
1113

1114

1115
def test_add_chunking_strategy_on_partition_auto_respects_max_chars():
1116
    filename = "example-docs/example-10k-1p.html"
1117

1118
    # default chunk size in chars is 200
1119
    partitioned_table_elements_200_chars = [
1120
        e
1121
        for e in partition(
1122
            filename,
1123
            chunking_strategy="by_title",
1124
            max_characters=200,
1125
            combine_text_under_n_chars=5,
1126
        )
1127
        if isinstance(e, (Table, TableChunk))
1128
    ]
1129

1130
    partitioned_table_elements_5_chars = [
1131
        e
1132
        for e in partition(
1133
            filename,
1134
            chunking_strategy="by_title",
1135
            max_characters=5,
1136
            combine_text_under_n_chars=5,
1137
        )
1138
        if isinstance(e, (Table, TableChunk))
1139
    ]
1140

1141
    elements = partition(filename)
1142

1143
    table_elements = [e for e in elements if isinstance(e, Table)]
1144

1145
    assert len(partitioned_table_elements_5_chars) != len(table_elements)
1146
    assert len(partitioned_table_elements_200_chars) != len(table_elements)
1147

1148
    # trailing whitespace is stripped from the first chunk, leaving only a checkbox character
1149
    assert len(partitioned_table_elements_5_chars[0].text) == 1
1150
    # but the second chunk is the full 5 characters
1151
    assert len(partitioned_table_elements_5_chars[1].text) == 5
1152
    assert len(partitioned_table_elements_5_chars[0].metadata.text_as_html) == 5
1153

1154
    # the first table element is under 200 chars so doesn't get chunked!
1155
    assert table_elements[0] == partitioned_table_elements_200_chars[0]
1156
    assert len(partitioned_table_elements_200_chars[0].text) < 200
1157
    assert len(partitioned_table_elements_200_chars[1].text) == 198
1158
    assert len(partitioned_table_elements_200_chars[1].metadata.text_as_html) == 200
1159

1160

1161
def test_add_chunking_strategy_chars_on_partition_auto_adds_is_continuation():
1162
    filename = "example-docs/example-10k-1p.html"
1163

1164
    table_elements = [e for e in partition(filename) if isinstance(e, Table)]
1165
    chunked_table_elements = [
1166
        e
1167
        for e in partition(
1168
            filename,
1169
            chunking_strategy="by_title",
1170
        )
1171
        if isinstance(e, Table)
1172
    ]
1173

1174
    assert table_elements != chunked_table_elements
1175

1176
    i = 0
1177
    for table in chunked_table_elements:
1178
        # have to reset the counter to 0 here when we encounter a Table element
1179
        if isinstance(table, Table):
1180
            i = 0
1181
        if i > 0 and isinstance(table, TableChunk):
1182
            assert table.metadata.is_continuation is True
1183
            i += 1
1184

1185

1186
EXAMPLE_LANG_DOCS = "example-docs/language-docs/eng_spa_mult."
1187

1188

1189
@pytest.mark.parametrize(
1190
    "file_extension",
1191
    [
1192
        "doc",
1193
        "docx",
1194
        "eml",
1195
        "epub",
1196
        "html",
1197
        "md",
1198
        "odt",
1199
        "org",
1200
        "ppt",
1201
        "pptx",
1202
        "rst",
1203
        "rtf",
1204
        "txt",
1205
        "xml",
1206
    ],
1207
)
1208
def test_partition_respects_language_arg(file_extension):
1209
    filename = EXAMPLE_LANG_DOCS + file_extension
1210
    elements = partition(filename=filename, languages=["deu"])
1211
    assert all(element.metadata.languages == ["deu"] for element in elements)
1212

1213

1214
def test_partition_respects_detect_language_per_element_arg():
1215
    filename = "example-docs/language-docs/eng_spa_mult.txt"
1216
    elements = partition(filename=filename, detect_language_per_element=True)
1217
    langs = [element.metadata.languages for element in elements]
1218
    assert langs == [["eng"], ["spa", "eng"], ["eng"], ["eng"], ["spa"]]
1219

1220

1221
# check that the ["eng"] default in `partition` does not overwrite the ["auto"]
1222
# default in other `partition_` functions.
1223
def test_partition_default_does_not_overwrite_other_defaults():
1224
    # the default for `languages` is ["auto"] in partiton_text
1225
    from unstructured.partition.text import partition_text
1226

1227
    # Use a document that is primarily in a language other than English
1228
    filename = "example-docs/language-docs/UDHR_first_article_all.txt"
1229
    text_elements = partition_text(filename)
1230
    assert text_elements[0].metadata.languages != ["eng"]
1231

1232
    auto_elements = partition(filename)
1233
    assert auto_elements[0].metadata.languages != ["eng"]
1234
    assert auto_elements[0].metadata.languages == text_elements[0].metadata.languages
1235

1236

1237
def test_partition_languages_default_to_None():
1238
    filename = "example-docs/handbook-1p.docx"
1239
    elements = partition(filename=filename, detect_language_per_element=True)
1240
    # PageBreak and other elements with no text will have `None` for `languages`
1241
    none_langs = [element for element in elements if element.metadata.languages is None]
1242
    assert none_langs[0].text == ""
1243

1244

1245
def test_partition_languages_incorrectly_defaults_to_English(tmpdir):
1246
    # We don't totally rely on langdetect for short text, so text like the following that is
1247
    # in German will be labeled as English.
1248
    german = "Ein kurzer Satz."
1249
    filepath = os.path.join(tmpdir, "short-german.txt")
1250
    with open(filepath, "w") as f:
1251
        f.write(german)
1252
    elements = partition(filepath)
1253
    assert elements[0].metadata.languages == ["eng"]
1254

1255

1256
def test_partition_timeout_gets_routed():
1257
    class CallException(Exception):
1258
        pass
1259

1260
    mock_ocr_func = Mock(side_effect=CallException("Function called!"))
1261
    with patch("unstructured.partition.auto.file_and_type_from_url", mock_ocr_func), pytest.raises(
1262
        CallException
1263
    ):
1264
        auto.partition(url="fake_url", request_timeout=326)
1265
    kwargs = mock_ocr_func.call_args.kwargs
1266
    assert "request_timeout" in kwargs
1267
    assert kwargs["request_timeout"] == 326
1268

1269

1270
def test_partition_image_with_bmp_with_auto(
1271
    tmpdir,
1272
    filename="example-docs/layout-parser-paper-with-table.jpg",
1273
):
1274
    bmp_filename = os.path.join(tmpdir.dirname, "example.bmp")
1275
    img = Image.open(filename)
1276
    img.save(bmp_filename)
1277

1278
    elements = partition(
1279
        filename=bmp_filename,
1280
        strategy=PartitionStrategy.HI_RES,
1281
    )
1282
    table = [el.metadata.text_as_html for el in elements if el.metadata.text_as_html]
1283
    assert len(table) == 1
1284
    assert "<table><thead><th>" in table[0]
1285

1286

1287
def test_auto_partition_eml_add_signature_to_metadata():
1288
    elements = partition(filename="example-docs/eml/signed-doc.p7s")
1289
    assert len(elements) == 1
1290
    assert elements[0].text == "This is a test"
1291
    assert elements[0].metadata.signature == "<SIGNATURE>\n"
1292

Использование cookies

Мы используем файлы cookie в соответствии с Политикой конфиденциальности и Политикой использования cookies.

Нажимая кнопку «Принимаю», Вы даете АО «СберТех» согласие на обработку Ваших персональных данных в целях совершенствования нашего веб-сайта и Сервиса GitVerse, а также повышения удобства их использования.

Запретить использование cookies Вы можете самостоятельно в настройках Вашего браузера.