diff --git a/CHANGELOG.md b/CHANGELOG.md index 5f1d11f8f6..bf7f83db11 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,9 @@ +## 0.25.2 + +### Fixes + +- **Avoid copying spooled uploads into memory**: DOCX, PPTX, and shared partitioning paths now reuse `SpooledTemporaryFile` inputs directly on supported Python versions instead of copying their complete contents into `BytesIO`. Large uploads remain disk-backed, avoiding an additional document-sized heap allocation without changing the partition API. + ## 0.25.1 ### Fixes diff --git a/test_unstructured/partition/common/test_common.py b/test_unstructured/partition/common/test_common.py index 35825dedf4..fa480f52c9 100644 --- a/test_unstructured/partition/common/test_common.py +++ b/test_unstructured/partition/common/test_common.py @@ -1,5 +1,6 @@ import pathlib from multiprocessing import Pool +from tempfile import SpooledTemporaryFile import numpy as np import pytest @@ -29,6 +30,17 @@ from unstructured.partition.common import common +def test_spooled_to_bytes_io_if_needed_rewinds_without_copying(): + with SpooledTemporaryFile(max_size=1, mode="w+b") as spooled_file: + spooled_file.write(b"sample content") + + result = common.spooled_to_bytes_io_if_needed(spooled_file) + + assert result is spooled_file + assert result.tell() == 0 + assert result.read() == b"sample content" + + class MockPageLayout(layout.PageLayout): def __init__(self, number: int, image: Image.Image): self.number = number diff --git a/test_unstructured/partition/test_docx.py b/test_unstructured/partition/test_docx.py index 6911fefc52..6e5cd67bb3 100644 --- a/test_unstructured/partition/test_docx.py +++ b/test_unstructured/partition/test_docx.py @@ -919,9 +919,7 @@ def it_uses_the_path_to_open_the_presentation_when_file_path_is_provided( assert opts._docx_file == "l/m/n.docx" - def and_it_uses_a_BytesIO_file_to_replaces_a_SpooledTemporaryFile_provided( - self, opts_args: dict[str, Any] - ): + def and_it_uses_a_SpooledTemporaryFile_directly(self, opts_args: dict[str, Any]): with tempfile.SpooledTemporaryFile() as spooled_temp_file: spooled_temp_file.write(b"abcdefg") opts_args["file"] = spooled_temp_file @@ -929,9 +927,8 @@ def and_it_uses_a_BytesIO_file_to_replaces_a_SpooledTemporaryFile_provided( docx_file = opts._docx_file - assert docx_file is not spooled_temp_file - assert isinstance(docx_file, io.BytesIO) - assert docx_file.getvalue() == b"abcdefg" + assert docx_file is spooled_temp_file + assert docx_file.read() == b"abcdefg" def and_it_uses_the_provided_file_directly_when_not_a_SpooledTemporaryFile( self, opts_args: dict[str, Any] diff --git a/test_unstructured/partition/test_pptx.py b/test_unstructured/partition/test_pptx.py index d47d4a6bc1..3e63109a74 100644 --- a/test_unstructured/partition/test_pptx.py +++ b/test_unstructured/partition/test_pptx.py @@ -700,9 +700,7 @@ def it_uses_the_path_to_open_the_presentation_when_file_path_is_provided( assert opts.pptx_file == "l/m/n.pptx" - def and_it_uses_a_BytesIO_file_to_replaces_a_SpooledTemporaryFile_provided( - self, opts_args: dict[str, Any] - ): + def and_it_uses_a_SpooledTemporaryFile_directly(self, opts_args: dict[str, Any]): with tempfile.SpooledTemporaryFile() as spooled_temp_file: spooled_temp_file.write(b"abcdefg") opts_args["file"] = spooled_temp_file @@ -710,9 +708,8 @@ def and_it_uses_a_BytesIO_file_to_replaces_a_SpooledTemporaryFile_provided( pptx_file = opts.pptx_file - assert pptx_file is not spooled_temp_file - assert isinstance(pptx_file, io.BytesIO) - assert pptx_file.getvalue() == b"abcdefg" + assert pptx_file is spooled_temp_file + assert pptx_file.read() == b"abcdefg" def and_it_uses_the_provided_file_directly_when_not_a_SpooledTemporaryFile( self, opts_args: dict[str, Any] diff --git a/unstructured/__version__.py b/unstructured/__version__.py index c7cc966b51..b20e30fc9b 100644 --- a/unstructured/__version__.py +++ b/unstructured/__version__.py @@ -1 +1 @@ -__version__ = "0.25.1" # pragma: no cover +__version__ = "0.25.2" # pragma: no cover diff --git a/unstructured/partition/common/common.py b/unstructured/partition/common/common.py index d18fc8c87b..91e45ada65 100644 --- a/unstructured/partition/common/common.py +++ b/unstructured/partition/common/common.py @@ -6,7 +6,7 @@ from io import BufferedReader, BytesIO, TextIOWrapper from tempfile import SpooledTemporaryFile from time import sleep -from typing import IO, TYPE_CHECKING, Any, Optional, TypeVar, cast +from typing import IO, TYPE_CHECKING, Any, Optional, TypeVar import emoji import psutil @@ -350,19 +350,20 @@ def exactly_one(**kwargs: Any) -> None: _T = TypeVar("_T") -def spooled_to_bytes_io_if_needed(file: _T | SpooledTemporaryFile[bytes]) -> _T | BytesIO: - """Convert `file` to `BytesIO` when it is a `SpooledTemporaryFile`. +def spooled_to_bytes_io_if_needed(file: _T) -> _T: + """Rewind and return a `SpooledTemporaryFile` without copying its contents. Note that `file` does not need to be IO[bytes]. It can be `None` or `bytes` and this function will not complain. - In Python <3.11, `SpooledTemporaryFile` does not implement `.readable()` or `.seekable()` which - triggers an exception when the file is loaded by certain packages. In particular, the stdlib - `zipfile.Zipfile` raises on opening a `SpooledTemporaryFile` as does `Pandas.read_csv()`. + Python 3.11 and newer provide the complete buffered-I/O interface required by consumers such + as `zipfile.ZipFile` and `pandas.read_csv()`. Since those are the only Python versions this + package supports, converting the spool to `BytesIO` only adds a document-sized allocation. + + The function name is retained for compatibility with existing call sites. """ if isinstance(file, SpooledTemporaryFile): file.seek(0) - return BytesIO(cast(bytes, file.read())) # -- return `file` unchanged otherwise -- return file diff --git a/unstructured/partition/docx.py b/unstructured/partition/docx.py index 1347d88134..9301fbcd09 100644 --- a/unstructured/partition/docx.py +++ b/unstructured/partition/docx.py @@ -2,7 +2,6 @@ from __future__ import annotations -import io import itertools import logging import os @@ -341,12 +340,8 @@ def _docx_file(self) -> str | IO[bytes]: if self._file_path: return self._file_path - # -- In Python <3.11 SpooledTemporaryFile does not implement ".seekable" which triggers an - # -- exception when Zipfile tries to open it. The docx format is a zip archive so we need - # -- to work around that bug here. if isinstance(self._file, tempfile.SpooledTemporaryFile): self._file.seek(0) - return io.BytesIO(self._file.read()) assert self._file is not None # -- assured by `._validate()` -- return self._file diff --git a/unstructured/partition/pptx.py b/unstructured/partition/pptx.py index 8652234255..1e5da3753a 100644 --- a/unstructured/partition/pptx.py +++ b/unstructured/partition/pptx.py @@ -6,7 +6,6 @@ from __future__ import annotations -import io from functools import cached_property from tempfile import SpooledTemporaryFile from typing import IO, Any, Iterator, Protocol, Sequence @@ -447,14 +446,9 @@ def pptx_file(self) -> str | IO[bytes]: if self._file_path: return self._file_path - # -- In Python <3.11 SpooledTemporaryFile does not implement ".seekable" which triggers an - # -- exception when Zipfile tries to open it. The pptx format is a zip archive so we need - # -- to work around that bug here. - if isinstance(self._file, SpooledTemporaryFile): - self._file.seek(0) - return io.BytesIO(self._file.read()) - if self._file: + if isinstance(self._file, SpooledTemporaryFile): + self._file.seek(0) return self._file raise ValueError(