diff --git a/test_unstructured/partition/common/test_common.py b/test_unstructured/partition/common/test_common.py index 35825dedf4..b12b1e435d 100644 --- a/test_unstructured/partition/common/test_common.py +++ b/test_unstructured/partition/common/test_common.py @@ -2,6 +2,8 @@ from multiprocessing import Pool import numpy as np +from io import BytesIO +from unstructured.partition.common.common import convert_to_bytes import pytest from PIL import Image from unstructured_inference.constants import IsExtracted @@ -466,3 +468,10 @@ def test_normalize_layout_element_layout_element_text_source_metadata(): assert hasattr(element, "metadata") assert hasattr(element.metadata, "is_extracted") assert element.metadata.is_extracted == "true" + +def test_convert_to_bytes_preserves_seekable_file_position(): + file = BytesIO(b"hello world") + file.seek(6) + + assert convert_to_bytes(file) == b"hello world" + assert file.tell() == 6 \ No newline at end of file diff --git a/test_unstructured/partition/test_auto.py b/test_unstructured/partition/test_auto.py index d7927ccc92..42802e2dd3 100644 --- a/test_unstructured/partition/test_auto.py +++ b/test_unstructured/partition/test_auto.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import zipfile import os import pathlib import tempfile @@ -1118,6 +1119,28 @@ def test_auto_partition_forwards_include_page_breaks_to_partition_pdf(): # -- metadata_filename ---------------------------------------------------- +def test_partition_txt_from_zip_ext_file(tmp_path): + zip_path = tmp_path / "text_file.zip" + + with zipfile.ZipFile(zip_path, "w") as zf: + zf.writestr( + "hello.txt", + "Hello from inside a zip file.\nThis should be partitioned successfully.", + ) + + with zipfile.ZipFile(zip_path, "r") as zf: + info = zf.infolist()[0] + + with zf.open(info) as file: + elements = partition( + file=file, + metadata_filename=info.filename, + ) + + assert [element.text for element in elements] == [ + "Hello from inside a zip file.", + "This should be partitioned successfully.", + ] def test_auto_partition_forwards_metadata_filename_via_kwargs(): with open(example_doc_path("fake-text.txt"), "rb") as f: diff --git a/unstructured/partition/common/common.py b/unstructured/partition/common/common.py index d18fc8c87b..56b3c8850c 100644 --- a/unstructured/partition/common/common.py +++ b/unstructured/partition/common/common.py @@ -371,7 +371,8 @@ def spooled_to_bytes_io_if_needed(file: _T | SpooledTemporaryFile[bytes]) -> _T def convert_to_bytes(file: bytes | IO[bytes]) -> bytes: """Extract the bytes from `file` without preventing it from being read again later. - As a convenience to simplify client code, also returns `file` unchanged if it is already bytes. + As a convenience to simplify client code, also returns `file` unchanged if it is already + bytes. """ if isinstance(file, bytes): return file @@ -389,6 +390,30 @@ def convert_to_bytes(file: bytes | IO[bytes]) -> bytes: with open(file.name, "rb") as f: return f.read() + if hasattr(file, "read"): + original_position = None + + if hasattr(file, "tell") and hasattr(file, "seek"): + try: + original_position = file.tell() + file.seek(0) + except (OSError, ValueError): + original_position = None + + data = file.read() + + if original_position is not None: + try: + file.seek(original_position) + except (OSError, ValueError): + pass + + if isinstance(data, str): + return data.encode() + + if isinstance(data, bytes): + return data + raise ValueError("Invalid file-like object type")