diff --git a/CHANGELOG.md b/CHANGELOG.md index b75a37ce..7e1148cd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,9 @@ +## 0.1.9 + +### Performance + +- **Avoid copying uploads during MIME detection**: unknown file types are now detected directly from the existing spooled upload stream while preserving the filename expected by the file-type detector. + ## 0.1.8 ### Features diff --git a/prepline_general/api/__version__.py b/prepline_general/api/__version__.py index 93b52aec..4eeb8fc3 100644 --- a/prepline_general/api/__version__.py +++ b/prepline_general/api/__version__.py @@ -1 +1 @@ -__version__ = "0.1.8" # pragma: no cover +__version__ = "0.1.9" # pragma: no cover diff --git a/prepline_general/api/filetypes.py b/prepline_general/api/filetypes.py index c010afaa..062cff22 100644 --- a/prepline_general/api/filetypes.py +++ b/prepline_general/api/filetypes.py @@ -1,6 +1,4 @@ -import os -from typing import Optional -from io import BytesIO +from typing import IO, Any, Optional, cast from fastapi import HTTPException, UploadFile @@ -8,6 +6,26 @@ from unstructured.file_utils.model import FileType +class _SpooledFileProxy: + """Expose a spooled upload as a normal file without copying its complete body. + + `unstructured.detect_filetype()` copies `SpooledTemporaryFile` inputs into a `BytesIO`, even + though its detector only seeks and reads small portions. Hiding the concrete type preserves the + normal file interface while avoiding that document-sized allocation. + + `name` is set explicitly rather than forwarded: the detector reads `.name` to derive the + filename-extension, and a spooled file's own `.name` is the temporary file rather than the + uploaded filename. + """ + + def __init__(self, file: IO[bytes], name: str | None): + self._file = file + self.name = name + + def __getattr__(self, name: str) -> Any: + return getattr(self._file, name) + + def _remove_optional_info_from_mime_type(content_type: str | None) -> str | None: """removes charset information from mime types, e.g., "application/json; charset=utf-8" -> "application/json" @@ -37,14 +55,14 @@ def get_validated_mimetype(file: UploadFile, content_type_hint: str | None = Non filetype = FileType.from_mime_type(content_type) # If content_type was not specified, use the library to identify the file - # We inspect the bytes to do this, so we need to buffer the file + # The detector seeks and reads bounded portions of the upload. Proxy the spooled file so the + # dependency does not make a complete in-memory `BytesIO` copy first. if not filetype or filetype == FileType.UNK: - file_buffer = BytesIO(file.file.read()) - file.file.seek(0) - - file_buffer.name = file.filename - - filetype = detect_filetype(file=file_buffer) + file_proxy = cast(IO[bytes], _SpooledFileProxy(file.file, file.filename)) + try: + filetype = detect_filetype(file=file_proxy) + finally: + file.file.seek(0) if not filetype.is_partitionable: raise HTTPException( diff --git a/test_general/api/test_filetypes.py b/test_general/api/test_filetypes.py new file mode 100644 index 00000000..d25b2d42 --- /dev/null +++ b/test_general/api/test_filetypes.py @@ -0,0 +1,43 @@ +from tempfile import SpooledTemporaryFile + +import pytest +from fastapi import UploadFile + +from prepline_general.api import filetypes +from unstructured.file_utils.model import FileType + + +def test_unknown_mimetype_is_detected_from_existing_upload_stream(monkeypatch): + upload_stream = SpooledTemporaryFile() + upload_stream.write(b"sample text") + upload_stream.seek(0) + upload = UploadFile(file=upload_stream, filename="sample.txt") + + def fake_detect_filetype(*, file): + assert file._file is upload_stream + assert file.name == "sample.txt" + file.seek(4) + return FileType.TXT + + monkeypatch.setattr(filetypes, "detect_filetype", fake_detect_filetype) + + assert filetypes.get_validated_mimetype(upload) == "text/plain" + assert upload_stream.tell() == 0 + + +def test_unknown_mimetype_rewinds_upload_stream_when_detection_fails(monkeypatch): + upload_stream = SpooledTemporaryFile() + upload_stream.write(b"sample text") + upload_stream.seek(0) + upload = UploadFile(file=upload_stream, filename="sample.txt") + + def fake_detect_filetype(*, file): + file.seek(4) + raise RuntimeError("detection failed") + + monkeypatch.setattr(filetypes, "detect_filetype", fake_detect_filetype) + + with pytest.raises(RuntimeError, match="detection failed"): + filetypes.get_validated_mimetype(upload) + + assert upload_stream.tell() == 0