Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
File size: 3,177 Bytes
3530326 3719fca 3530326 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | from __future__ import annotations
from pathlib import Path
import pypdfium2 as pdfium
import pytest
from PIL import Image
from document_processing import (
MAX_DOCUMENT_BYTES,
MAX_IMAGE_EDGE,
MAX_PAGES,
DocumentError,
detect_document_kind,
prepare_document,
)
def make_image(path: Path, *, size: tuple[int, int] = (320, 200)) -> None:
image = Image.new("RGB", size, "#ffffff")
image.save(path, format="PNG")
def make_pdf(path: Path, pages: int) -> None:
document = pdfium.PdfDocument.new()
try:
for _ in range(pages):
page = document.new_page(width=612, height=792)
page.close()
document.save(path)
finally:
document.close()
def test_detects_image_content_instead_of_extension(tmp_path: Path) -> None:
path = tmp_path / "not-really-a-pdf.pdf"
make_image(path)
assert detect_document_kind(path) == "png"
def test_reencodes_and_downsizes_image(tmp_path: Path) -> None:
path = tmp_path / "large.png"
make_image(path, size=(3000, 1000))
document = prepare_document(
path,
delete_source=False,
enforce_temp_location=False,
)
assert document.kind == "png"
assert document.total_pages == 1
assert document.pages[0].width == MAX_IMAGE_EDGE
assert document.pages[0].height < MAX_IMAGE_EDGE
assert document.pages[0].png.startswith(b"\x89PNG\r\n\x1a\n")
def test_caps_pdf_at_ten_pages(tmp_path: Path) -> None:
path = tmp_path / "eleven-pages.pdf"
make_pdf(path, MAX_PAGES + 1)
document = prepare_document(
path,
delete_source=False,
enforce_temp_location=False,
)
assert document.total_pages == MAX_PAGES + 1
assert document.processed_pages == MAX_PAGES
assert [page.number for page in document.pages] == list(
range(1, MAX_PAGES + 1)
)
def test_deletes_source_after_preparation(tmp_path: Path) -> None:
path = tmp_path / "upload.png"
make_image(path)
prepare_document(
path,
delete_source=True,
enforce_temp_location=False,
)
assert not path.exists()
def test_rejects_invalid_and_oversized_documents(tmp_path: Path) -> None:
invalid = tmp_path / "invalid.pdf"
invalid.write_bytes(b"not a document")
with pytest.raises(DocumentError, match="valid PDF"):
prepare_document(
invalid,
delete_source=False,
enforce_temp_location=False,
)
oversized = tmp_path / "oversized.pdf"
with oversized.open("wb") as handle:
handle.truncate(MAX_DOCUMENT_BYTES + 1)
with pytest.raises(DocumentError, match="20 MB"):
prepare_document(
oversized,
delete_source=False,
enforce_temp_location=False,
)
def test_rejects_non_temporary_path_when_enforced(tmp_path: Path) -> None:
path = tmp_path / "upload.png"
make_image(path)
# pytest's tmp_path is already below the system temp directory and is valid.
document = prepare_document(
path,
delete_source=False,
enforce_temp_location=True,
)
assert document.processed_pages == 1
|