|
48 | 48 | from unstructured.file_utils.model import FileType, create_file_type, register_partitioner |
49 | 49 | from unstructured.partition.auto import _PartitionerLoader, partition |
50 | 50 | from unstructured.partition.common import UnsupportedFileFormatError |
| 51 | +from unstructured.partition.common.metadata import is_attachment_element |
51 | 52 | from unstructured.partition.utils.constants import PartitionStrategy |
52 | 53 | from unstructured.staging.base import elements_from_json, elements_to_dicts, elements_to_json |
53 | 54 |
|
@@ -1417,6 +1418,46 @@ def test_auto_partition_applies_the_correct_filetype_for_all_filetypes( |
1417 | 1418 | ) |
1418 | 1419 |
|
1419 | 1420 |
|
| 1421 | +@pytest.mark.parametrize( |
| 1422 | + "file_name", |
| 1423 | + ["eml/fake-email-attachment.eml", "fake-email-attachment.msg"], |
| 1424 | +) |
| 1425 | +def test_auto_partition_preserves_the_filetype_of_attachment_elements(file_name: str): |
| 1426 | + """Attachment elements keep their own filetype, not the containing message's. |
| 1427 | +
|
| 1428 | + Their filetype was assigned by the nested `partition()` call that produced them, so the |
| 1429 | + outer call must not re-stamp them with e.g. `message/rfc822`. |
| 1430 | + """ |
| 1431 | + elements = partition(example_doc_path(file_name), process_attachments=True) |
| 1432 | + |
| 1433 | + attachment_elements = [e for e in elements if is_attachment_element(e)] |
| 1434 | + assert attachment_elements |
| 1435 | + assert all(e.metadata.filetype == FileType.TXT.mime_type for e in attachment_elements) |
| 1436 | + |
| 1437 | + |
| 1438 | +@pytest.mark.parametrize( |
| 1439 | + "file_name", |
| 1440 | + ["eml/fake-email-attachment.eml", "fake-email-attachment.msg"], |
| 1441 | +) |
| 1442 | +def test_auto_partition_preserves_attachment_filetype_when_container_filename_is_unknown( |
| 1443 | + file_name: str, |
| 1444 | +): |
| 1445 | + """The attachment guard must not depend on `.metadata.attached_to_filename`. |
| 1446 | +
|
| 1447 | + That field is `None` when the containing document's file-name is unknown -- partitioning a |
| 1448 | + file-like object with no `metadata_filename` -- so keying the guard on it would let the |
| 1449 | + containing document's filetype overwrite the attachment's in exactly that case. |
| 1450 | + """ |
| 1451 | + with open(example_doc_path(file_name), "rb") as f: |
| 1452 | + elements = partition(file=f, process_attachments=True) |
| 1453 | + |
| 1454 | + attachment_elements = [e for e in elements if is_attachment_element(e)] |
| 1455 | + assert attachment_elements |
| 1456 | + # -- precondition: this is the case the `attached_to_filename` guard cannot see -- |
| 1457 | + assert all(e.metadata.attached_to_filename is None for e in attachment_elements) |
| 1458 | + assert all(e.metadata.filetype == FileType.TXT.mime_type for e in attachment_elements) |
| 1459 | + |
| 1460 | + |
1420 | 1461 | def test_detect_filetype_maps_file_to_bytes_io_when_spooled_temp_file_used(mocker): |
1421 | 1462 | detect_filetype_mock = MagicMock(return_value=FileType.JSON) |
1422 | 1463 | mocker.patch("unstructured.file_utils.filetype._FileTypeDetector", detect_filetype_mock) |
|
0 commit comments