diff --git a/app/pybind_parse.cpp b/app/pybind_parse.cpp index 82d7367c..46dad203 100644 --- a/app/pybind_parse.cpp +++ b/app/pybind_parse.cpp @@ -908,7 +908,66 @@ PYBIND11_MODULE(pdf_parsers, m) { Returns: dict: A None or string of the metadata in xml of the document.)") - + + .def("get_attachments", + [](docling::docling_parser &self, const std::string &key) -> nlohmann::json { + return self.get_attachments(key); + }, + pybind11::arg("key"), + R"( + Retrieve attachments for the document identified by its unique key. + + Parameters: + key (str): The unique key of the document. + + Returns: + list: A JSON array of attachment metadata (name, mime_type, size, annotations). + + Raises: + RuntimeError: If the key is not loaded.)") + + .def("get_attachment_data", + [](docling::docling_parser &self, const std::string &key, int index, long long max_size) -> pybind11::bytes { + return self.get_attachment_data(key, index, max_size); + }, + pybind11::arg("key"), + pybind11::arg("index"), + pybind11::arg("max_size"), + R"( + Retrieve raw bytes for a single attachment. + + Parameters: + key (str): The unique key of the document. + index (int): Attachment index from get_attachments(). + max_size (int): Maximum allowed decoded size in bytes (required). + + Returns: + bytes: Decoded attachment payload. + + Raises: + RuntimeError: If the key is not loaded, index is out of range, + size exceeds max_size, or the stream cannot be decoded.)") + + .def("write_attachment_data", + [](docling::docling_parser &self, const std::string &key, int index, long long max_size, const std::string &path) { + self.write_attachment_data(key, index, max_size, path); + }, + pybind11::arg("key"), + pybind11::arg("index"), + pybind11::arg("max_size"), + pybind11::arg("path"), + R"( + Stream a single attachment directly to a file without buffering the full decoded payload. + + Parameters: + key (str): The unique key of the document. + index (int): Attachment index from get_attachments(). + max_size (int): Maximum allowed decoded size in bytes (required). + path (str): Destination file path. + + Raises: + RuntimeError: If size exceeds max_size or the stream cannot be decoded.)") + .def("get_page_decoder", [](docling::docling_parser &self, const std::string &key, diff --git a/docling_parse/pdf_parser.py b/docling_parse/pdf_parser.py index 5f351de3..6ec8dff1 100644 --- a/docling_parse/pdf_parser.py +++ b/docling_parse/pdf_parser.py @@ -3,10 +3,12 @@ import hashlib import logging import math +import os +import tempfile from enum import IntEnum from io import BytesIO from pathlib import Path -from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union +from typing import Any, BinaryIO, Dict, Iterator, List, Optional, Sequence, Tuple, Union from docling_core.types.doc.base import BoundingBox, CoordOrigin, ImageRefMode from docling_core.types.doc.document import ImageRef @@ -15,6 +17,8 @@ BoundingRectangle, ColorRGBA, Coord2D, + FileAttachmentAnnotation, + PdfAttachment, PdfHyperlink, PdfMetaData, PdfPageBoundaryType, @@ -77,6 +81,46 @@ _log = logging.getLogger(__name__) +class _AttachmentDeletingFile: + """File wrapper that unlinks its mkstemp path on close. + + Created once per large attachment; defined at module scope so the + class object is not re-created on every get_attachment_stream call. + """ + + def __init__(self, fd: int, path: str): + self._path = path + self._file = os.fdopen(fd, "w+b") + self.name = path # for test compat + + def read(self, *a, **kw): + return self._file.read(*a, **kw) + + def seek(self, *a, **kw): + return self._file.seek(*a, **kw) + + def tell(self): + return self._file.tell() + + def close(self): + try: + self._file.close() + finally: + try: + os.unlink(self._path) + except Exception: + pass + + def __getattr__(self, n): + return getattr(self._file, n) + + def __enter__(self): + return self + + def __exit__(self, *a): + self.close() + + class PdfTocEntry(BaseModel): """PDF table of contents entry (recursive structure). @@ -900,6 +944,144 @@ def get_annotations(self) -> PdfAnnotations | None: else: raise RuntimeError("This document is not loaded.") + def get_attachments(self) -> List[PdfAttachment]: + """Get attachment metadata (does not decode bytes). + + Returns: + List[PdfAttachment]: metadata with name, mime_type, size, and page annotations. + + Raises: + RuntimeError: if the document is not loaded (missing key). + """ + if not self.is_loaded(): + raise RuntimeError("This document is not loaded.") + + raw = self._parser.get_attachments(key=self._key) + # C++ get_attachments now throws on missing key, consistent with + # get_attachment_data — empty doc returns []. + result: List[PdfAttachment] = [] + for item in raw: + name = item.get("name", "attachment") + mime = item.get("mime_type") + size = int(item.get("size", 0) or 0) + anns: List[FileAttachmentAnnotation] = [] + for ja in item.get("annotations") or []: + try: + bbox_vals = ja.get("bbox") or [0, 0, 0, 0] + bbox = _to_bounding_rectangle(tuple(float(v) for v in bbox_vals)) # type: ignore[arg-type] + anns.append( + FileAttachmentAnnotation( + page_no=int(ja.get("page_no", 0)) + 1, bbox=bbox + ) + ) + except Exception: + _log.debug("Skipping malformed annotation %s", ja, exc_info=True) + result.append( + PdfAttachment(name=name, mime_type=mime, size=size, annotations=anns) + ) + return result + + def get_attachment_data(self, index: int, *, max_size: int) -> bytes: + """Get raw bytes for a single attachment (memory-safe, requires max_size). + + Args: + index: Attachment index from get_attachments(). + max_size: Maximum allowed decoded size (required, enforces caller policy). + + Raises: + ValueError/RuntimeError: if size exceeds max_size or decode fails. + """ + if not self.is_loaded(): + raise RuntimeError("This document is not loaded.") + if not isinstance(max_size, int): + raise TypeError("max_size is required and must be an int") + return bytes( + self._parser.get_attachment_data( + key=self._key, index=index, max_size=max_size + ) + ) + + # Attachments at or below this size are served from memory; larger ones + # are spooled to a temp file that deletes itself on close. + _ATTACHMENT_MEMORY_LIMIT_BYTES = 8 * 1024 * 1024 + + def get_attachment_stream(self, index: int, *, max_size: int) -> BinaryIO: + """Get a readable BinaryIO for one attachment (spills to temp file if large). + + Args: + index: Attachment index from get_attachments(). + max_size: Maximum allowed decoded size (required). + + Returns: + BinaryIO: for size <= 8 MB a BytesIO; otherwise an open file + positioned at 0 that deletes itself on close (via + mkstemp wrapper). Caller must close() when done. + + Raises: + ValueError/RuntimeError: if size exceeds max_size or decode fails. + """ + if not self.is_loaded(): + raise RuntimeError("This document is not loaded.") + if not isinstance(max_size, int): + raise TypeError("max_size is required and must be an int") + + # Use metadata to decide transport without decoding first. + # size==0 means unknown — fall back to max_size. + try: + atts = self.get_attachments() + meta_size = atts[index].size if 0 <= index < len(atts) else 0 + except Exception: + meta_size = 0 + effective_size = meta_size if meta_size > 0 else max_size + + if effective_size > max_size: + raise RuntimeError( + f"attachment size {effective_size} exceeds max_size {max_size}" + ) + + use_memory = effective_size <= self._ATTACHMENT_MEMORY_LIMIT_BYTES + # Also cap memory by max_size + if max_size <= self._ATTACHMENT_MEMORY_LIMIT_BYTES: + use_memory = True + + if use_memory: + data = self.get_attachment_data(index, max_size=max_size) + if len(data) <= self._ATTACHMENT_MEMORY_LIMIT_BYTES: + return BytesIO(data) + # Size was underestimated — fall through to file path + # but avoid double copy: spill via pipe if available + effective_size = len(data) + del data # free before file spill + + # Large path: stream directly to file without holding full bytes in memory. + # Prefer native pipe if available (avoids C++ Buffer allocation). + if hasattr(self._parser, "write_attachment_data"): + tmp_fd, tmp_path = tempfile.mkstemp(suffix="__attachment") + wrapper = _AttachmentDeletingFile(tmp_fd, tmp_path) + try: + self._parser.write_attachment_data( + key=self._key, index=index, max_size=max_size, path=tmp_path + ) + wrapper.seek(0) + return wrapper # type: ignore[return-value] + except Exception: + wrapper.close() + raise + else: + # Fallback: legacy path — still avoid 2x memory by deleting after write + data = self.get_attachment_data(index, max_size=max_size) + tmp = tempfile.NamedTemporaryFile(delete=True, suffix="__attachment") + try: + tmp.write(data) + tmp.flush() + tmp.seek(0) + return tmp + except Exception: + tmp.close() + raise + finally: + del data + def get_page( self, page_no: int, diff --git a/pyproject.toml b/pyproject.toml index bf20df2b..d4e9ddc5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ requires-python = ">=3.10" dependencies = [ "pillow>=10.0.0,<13.0.0", "pydantic>=2.0.0", - "docling-core>=2.85.0,<3.0.0", + "docling-core>=2.91.0,<3.0.0", "pywin32>=305; sys_platform == 'win32'", ] [project.urls] diff --git a/src/parse/pdf_decoders/document.h b/src/parse/pdf_decoders/document.h index b3e2386f..a06088cb 100644 --- a/src/parse/pdf_decoders/document.h +++ b/src/parse/pdf_decoders/document.h @@ -7,9 +7,11 @@ #include #include #include +#include #include #include +#include #include namespace pdflib @@ -35,6 +37,10 @@ namespace pdflib nlohmann::json get_meta_xml(); nlohmann::json get_table_of_contents(); + nlohmann::json get_attachments(); + std::shared_ptr get_attachment_data(int index, long long max_size); + void write_attachment_data(int index, long long max_size, const std::string& path); + bool process_document_from_file(std::string& _filename, std::optional& _password, bool keep_qpdf_warnings = false); @@ -73,6 +79,13 @@ namespace pdflib void ensure_annots_loaded(); + void ensure_attachments_loaded(); + + // Shared attachment helpers (deduped from get/write paths). + const AttachmentRecord& require_attachment_record(int index, long long max_size); + QPDFObjectHandle require_attachment_stream(int index, long long max_size); + static void warn_if_raw_length_exceeds(const QPDFObjectHandle& stream_obj, long long max_size); + void update_timings(pdf_timings& timings_, bool set_timer); private: @@ -92,6 +105,10 @@ namespace pdflib nlohmann::json json_annots; bool annots_loaded; + // Attachments (metadata only, lazy) + std::vector attachment_records; + bool attachments_loaded; + // New: Persistent page decoders for typed API std::map page_decoders; }; @@ -110,6 +127,8 @@ namespace pdflib json_annots(nlohmann::json::value_t::null), annots_loaded(false), + attachment_records({}), + attachments_loaded(false), page_decoders({}) { configure_qpdf_warnings(qpdf_document); @@ -129,6 +148,8 @@ namespace pdflib json_annots(nlohmann::json::value_t::null), annots_loaded(false), + attachment_records({}), + attachments_loaded(false), page_decoders({}) { configure_qpdf_warnings(qpdf_document); @@ -180,6 +201,145 @@ namespace pdflib return json_annots["table_of_contents"]; } + void pdf_decoder::ensure_attachments_loaded() + { + if(attachments_loaded) + return; + try + { + QPDFObjectHandle qpdf_root = qpdf_document.getRoot(); + attachment_records = extract_attachment_records(qpdf_document, qpdf_root); + } + catch(const std::exception& exc) + { + LOG_S(WARNING) << "filename: " << filename << " failed to extract attachments: " << exc.what(); + } + attachments_loaded = true; + } + + const AttachmentRecord& pdf_decoder::require_attachment_record(int index, long long max_size) + { + ensure_attachments_loaded(); + if(index < 0 || index >= static_cast(attachment_records.size())) + throw std::out_of_range("attachment index out of range"); + const auto& rec = attachment_records[static_cast(index)]; + if(rec.size > max_size) + throw std::runtime_error("attachment size " + std::to_string(rec.size) + " exceeds max_size " + std::to_string(max_size)); + if(!rec.obj_gen.isIndirect()) + throw std::runtime_error( + "attachment at index " + std::to_string(index) + + " has a direct (non-indirect) EF stream and cannot be fetched — direct streams have no indirect object ID"); + return rec; + } + + QPDFObjectHandle pdf_decoder::require_attachment_stream(int index, long long max_size) + { + const auto& rec = require_attachment_record(index, max_size); + QPDFObjectHandle stream_obj; + try { stream_obj = qpdf_document.getObjectByObjGen(rec.obj_gen); } + catch(const std::exception& exc) { LOG_S(WARNING) << "attachment stream lookup failed: " << exc.what(); } + if(!stream_obj.isStream()) + throw std::runtime_error("attachment stream not found for index " + std::to_string(index)); + return stream_obj; + } + + void pdf_decoder::warn_if_raw_length_exceeds(const QPDFObjectHandle& stream_obj, long long max_size) + { + // Guard against filter bombs: compare the declared raw /Length against the + // limit before decoding. The 10x slack only decides when a mismatch is + // worth a warning (compression can legitimately expand); the decoded size + // is still enforced separately. + // NB: stream handles don't proxy hasKey/getKey on this qpdf version — + // the dict must be addressed via getDict(). + try + { + QPDFObjectHandle dict = stream_obj.getDict(); + if(dict.hasKey("/Length") && dict.getKey("/Length").isInteger()) + { + long long raw_len = dict.getKey("/Length").getIntValue(); + if(raw_len > max_size * 10) + LOG_S(WARNING) << "attachment raw Length " << raw_len << " is much larger than max_size " << max_size; + } + } + catch(const std::exception& exc) + { + LOG_S(WARNING) << "failed to read attachment /Length: " << exc.what(); + } + } + + nlohmann::json pdf_decoder::get_attachments() + { + ensure_attachments_loaded(); + return attachments_to_json(attachment_records); + } + + std::shared_ptr pdf_decoder::get_attachment_data(int index, long long max_size) + { + QPDFObjectHandle stream_obj = require_attachment_stream(index, max_size); + warn_if_raw_length_exceeds(stream_obj, max_size); + + std::shared_ptr buf; + try + { + buf = stream_obj.getStreamData(qpdf_dl_all); + } + catch(const std::exception& exc) + { + throw std::runtime_error(std::string("failed to decode attachment stream: ") + exc.what()); + } + + if(!buf) + throw std::runtime_error("failed to get stream data for attachment " + std::to_string(index)); + + if(static_cast(buf->getSize()) > max_size) + throw std::runtime_error("decoded attachment size " + std::to_string(buf->getSize()) + " exceeds max_size " + std::to_string(max_size)); + + return buf; + } + + void pdf_decoder::write_attachment_data(int index, long long max_size, const std::string& path) + { + QPDFObjectHandle stream_obj = require_attachment_stream(index, max_size); + warn_if_raw_length_exceeds(stream_obj, max_size); + + std::ofstream out(path, std::ios::binary); + if(!out) + throw std::runtime_error("failed to open attachment output file: " + path); + + struct LimitedFilePipeline : public Pipeline + { + std::ofstream& out; + long long max_size; + long long written = 0; + LimitedFilePipeline(std::ofstream& o, long long m) : Pipeline("attachment-pipe", nullptr), out(o), max_size(m) {} + void write(unsigned char const* data, size_t len) override + { + written += static_cast(len); + if(written > max_size) + throw std::runtime_error("decoded attachment size " + std::to_string(written) + " exceeds max_size " + std::to_string(max_size)); + out.write(reinterpret_cast(data), len); + if(!out) + throw std::runtime_error("failed to write attachment data to file"); + } + void finish() override { out.flush(); } + }; + + LimitedFilePipeline pipeline(out, max_size); + try + { + if(!stream_obj.pipeStreamData(&pipeline, 0, qpdf_dl_all)) + throw std::runtime_error("failed to pipe attachment stream data for index " + std::to_string(index)); + pipeline.finish(); + } + catch(const std::exception&) + { + throw; + } + + if(!out) + throw std::runtime_error("failed to write attachment data to file: " + path); + } + bool pdf_decoder::process_document_from_file(std::string& _filename, std::optional& _password, bool keep_qpdf_warnings) @@ -228,6 +388,12 @@ namespace pdflib buffer = _buffer; password = _password; + // Reset cached state + attachment_records.clear(); + attachments_loaded = false; + annots_loaded = false; + json_annots = nlohmann::json::value_t::null; + LOG_S(INFO) << "start processing buffer of size " << buffer->size() << " by qpdf ..."; utils::timer timer; diff --git a/src/parse/qpdf/attachments.h b/src/parse/qpdf/attachments.h new file mode 100644 index 00000000..8256ab65 --- /dev/null +++ b/src/parse/qpdf/attachments.h @@ -0,0 +1,336 @@ +#ifndef QPDF_ATTACHMENTS_H +#define QPDF_ATTACHMENTS_H + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include + +namespace pdflib +{ + +struct AttachmentAnnotation +{ + int page_no = -1; // 0-based + std::array bbox = {0, 0, 0, 0}; +}; + +struct AttachmentRecord +{ + std::string name; + std::string mime_type; // empty if missing + long long size = 0; + // QPDF object identity of the EF stream. A non-indirect (0,0) value is the + // sentinel for a direct (non-indirect) EF stream: the stream dict is + // embedded inline in the FileSpec and has no indirect object ID, so it + // cannot be re-fetched via getObjectByID / getObjectByObjGen. + // Direct streams are intentionally distinct per dedup logic and are not + // cached; get_attachment_data() detects this sentinel and throws. + // An indirect stream genuinely at object 0 is not possible in a valid PDF + // (object 0 is the free-list head), so 0 is unambiguous here. + QPDFObjGen obj_gen; + std::vector annotations; +}; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +inline QPDFObjectHandle getStreamDict(QPDFObjectHandle oh) +{ + return oh.isStream() ? oh.getDict() : oh; +} + +inline std::string extract_attachment_name(QPDFObjectHandle fs) +{ + try + { + if(fs.hasKey("/UF") && fs.getKey("/UF").isString()) + { + std::string v = fs.getKey("/UF").getUTF8Value(); + if(!v.empty()) + return utils::string::fix_into_valid_utf8(v); + } + if(fs.hasKey("/F") && fs.getKey("/F").isString()) + { + std::string v = fs.getKey("/F").getUTF8Value(); + if(!v.empty()) + return utils::string::fix_into_valid_utf8(v); + } + } + catch(const std::exception& exc) + { + LOG_S(WARNING) << "failed to extract attachment name: " << exc.what(); + } + return "attachment"; +} + +inline std::string extract_attachment_mime(QPDFObjectHandle stream_obj) +{ + try + { + QPDFObjectHandle dict = getStreamDict(stream_obj); + if(dict.hasKey("/Subtype") && dict.getKey("/Subtype").isName()) + { + std::string m = dict.getKey("/Subtype").getName(); + if(!m.empty() && m[0]=='/') + m = m.substr(1); + return utils::string::fix_into_valid_utf8(m); + } + } + catch(const std::exception& exc) + { + LOG_S(WARNING) << "failed to extract mime: " << exc.what(); + } + return ""; +} + +inline long long extract_attachment_size(QPDFObjectHandle stream_obj) +{ + try + { + QPDFObjectHandle dict = getStreamDict(stream_obj); + if(dict.hasKey("/Params") && dict.getKey("/Params").isDictionary()) + { + QPDFObjectHandle params = dict.getKey("/Params"); + if(params.hasKey("/Size") && params.getKey("/Size").isInteger()) + return params.getKey("/Size").getIntValue(); + } + if(dict.hasKey("/DL") && dict.getKey("/DL").isInteger()) + return dict.getKey("/DL").getIntValue(); + if(dict.hasKey("/Length") && dict.getKey("/Length").isInteger()) + return dict.getKey("/Length").getIntValue(); + } + catch(const std::exception& exc) + { + LOG_S(WARNING) << "failed to extract size: " << exc.what(); + } + return 0; +} + +inline bool extract_rect(QPDFObjectHandle rect, std::array& out) +{ + if(!rect.isArray() || rect.getArrayNItems()!=4) + return false; + try + { + for(int i=0;i<4;++i) + out[i] = rect.getArrayItem(i).getNumericValue(); + return true; + } + catch(const std::exception& exc) + { + LOG_S(WARNING) << "malformed Rect: " << exc.what(); + return false; + } +} + +// One annotation entry for an attachment; bbox zeroed when /Rect is malformed. +inline AttachmentAnnotation make_attachment_annotation(int page_no, QPDFObjectHandle rect) +{ + AttachmentAnnotation ann; + ann.page_no = page_no; + if(!extract_rect(rect, ann.bbox)) + ann.bbox = {0,0,0,0}; + return ann; +} + +// --------------------------------------------------------------------------- +// Core extraction — metadata only, never calls getStreamData() +// --------------------------------------------------------------------------- + +inline std::vector extract_attachment_records(QPDF& pdf, QPDFObjectHandle& root) +{ + std::vector records; + std::map obj_to_index; + + auto process_filespec = [&](QPDFObjectHandle fs, int page_no, QPDFObjectHandle rect) { + if(!fs.isDictionary()) + { + LOG_S(WARNING) << "FileSpec is not a dict, skipping"; + return; + } + + std::string name = extract_attachment_name(fs); + + if(!fs.hasKey("/EF") || !fs.getKey("/EF").isDictionary()) + { + LOG_S(WARNING) << "FileSpec missing /EF for '" << name << "', skipping"; + return; + } + + QPDFObjectHandle ef_dict = fs.getKey("/EF"); + QPDFObjectHandle stream_obj; + + if(ef_dict.hasKey("/UF") && ef_dict.getKey("/UF").isStream()) + stream_obj = ef_dict.getKey("/UF"); + else if(ef_dict.hasKey("/F") && ef_dict.getKey("/F").isStream()) + stream_obj = ef_dict.getKey("/F"); + else + { + LOG_S(WARNING) << "EF dict has no stream for '" << name << "', skipping"; + return; + } + + if(!stream_obj.isStream()) + { + LOG_S(WARNING) << "EF entry is not a stream for '" << name << "'"; + return; + } + + QPDFObjGen og = stream_obj.getObjGen(); + // Direct (non-indirect) EF streams are embedded inline in the FileSpec + // dict. Each occurrence is a distinct object by definition, so no dedup + // is attempted for them. Only indirect streams have a stable object + // identity that can be shared across FileSpecs/annotations. + bool is_dedupable = og.isIndirect(); + + std::string mime = extract_attachment_mime(stream_obj); + long long size = extract_attachment_size(stream_obj); + + if(is_dedupable && obj_to_index.count(og)) + { + size_t idx = obj_to_index[og]; + if(page_no>=0) + records[idx].annotations.push_back(make_attachment_annotation(page_no, rect)); + return; + } + + AttachmentRecord rec; + rec.name = name; + rec.mime_type = mime; + rec.size = size; + rec.obj_gen = og; + + if(page_no>=0) + rec.annotations.push_back(make_attachment_annotation(page_no, rect)); + + records.push_back(std::move(rec)); + if(is_dedupable) + obj_to_index[og] = records.size()-1; + }; + + try + { + if(root.hasKey("/Names") && root.getKey("/Names").isDictionary()) + { + QPDFObjectHandle names = root.getKey("/Names"); + if(names.hasKey("/EmbeddedFiles") && names.getKey("/EmbeddedFiles").isDictionary()) + { + QPDFObjectHandle ef_tree = names.getKey("/EmbeddedFiles"); + std::function walk = [&](QPDFObjectHandle node){ + if(!node.isDictionary()) return; + if(node.hasKey("/Names") && node.getKey("/Names").isArray()) + { + QPDFObjectHandle arr = node.getKey("/Names"); + int n = arr.getArrayNItems(); + for(int i=1;i pages = pdf.getAllPages(); + for(size_t p=0; p(p); + if(!page.hasKey("/Annots")) + continue; + QPDFObjectHandle annots = page.getKey("/Annots"); + if(!annots.isArray()) + continue; + for(int i=0;i& records) +{ + nlohmann::json arr = nlohmann::json::array(); + for(const auto& r : records) + { + nlohmann::json item; + item["name"] = r.name; + item["mime_type"] = r.mime_type.empty() ? nlohmann::json(nullptr) : nlohmann::json(r.mime_type); + item["size"] = r.size; + nlohmann::json anns = nlohmann::json::array(); + for(const auto& a : r.annotations) + { + nlohmann::json ja; + ja["page_no"] = a.page_no; + ja["bbox"] = nlohmann::json::array({a.bbox[0], a.bbox[1], a.bbox[2], a.bbox[3]}); + anns.push_back(ja); + } + item["annotations"] = anns; + arr.push_back(std::move(item)); + } + return arr; +} + +} // namespace pdflib + +#endif diff --git a/src/pybind/docling_parser.h b/src/pybind/docling_parser.h index 8d82ec73..f587efe2 100644 --- a/src/pybind/docling_parser.h +++ b/src/pybind/docling_parser.h @@ -57,6 +57,9 @@ namespace docling nlohmann::json get_meta_xml(std::string key); nlohmann::json get_table_of_contents(std::string key); + nlohmann::json get_attachments(std::string key); + pybind11::bytes get_attachment_data(std::string key, int index, long long max_size); + void write_attachment_data(std::string key, int index, long long max_size, std::string path); std::shared_ptr> get_page_decoder(std::string key, int page, @@ -429,6 +432,35 @@ namespace docling return (itr->second)->get_table_of_contents(); } + nlohmann::json docling_parser::get_attachments(std::string key) + { + LOG_S(INFO) << __FUNCTION__; + + auto itr = doc_decoders.find(key); + if(itr==doc_decoders.end()) + throw std::runtime_error("key not found: " + key); + + return (itr->second)->get_attachments(); + } + + pybind11::bytes docling_parser::get_attachment_data(std::string key, int index, long long max_size) + { + auto itr = doc_decoders.find(key); + if(itr==doc_decoders.end()) + throw std::runtime_error("key not found: " + key); + + auto buf = (itr->second)->get_attachment_data(index, max_size); + return pybind11::bytes(reinterpret_cast(buf->getBuffer()), buf->getSize()); + } + + void docling_parser::write_attachment_data(std::string key, int index, long long max_size, std::string path) + { + auto itr = doc_decoders.find(key); + if(itr==doc_decoders.end()) + throw std::runtime_error("key not found: " + key); + (itr->second)->write_attachment_data(index, max_size, path); + } + std::shared_ptr> docling_parser::get_page_decoder(std::string key, int page, const pdflib::decode_config& config) diff --git a/tests/test_attachments.py b/tests/test_attachments.py new file mode 100644 index 00000000..a009efe3 --- /dev/null +++ b/tests/test_attachments.py @@ -0,0 +1,276 @@ +"""Attachment extraction tests — streaming, memory-safe API.""" + +import os +import tempfile +from io import BytesIO + +import pytest + +from docling_parse.pdf_parser import DoclingPdfParser + + +def _build_pdf(objects: dict[int, bytes]) -> bytes: + """Build a minimal PDF from obj_num -> raw_content bytes.""" + header = b"%PDF-1.4\n%\xe2\xe3\xcf\xd3\n" + parts: list[bytes] = [header] + offsets: dict[int, int] = {} + for num in sorted(objects.keys()): + offsets[num] = sum(len(p) for p in parts) + parts.append(f"{num} 0 obj\n".encode()) + parts.append(objects[num]) + if not objects[num].endswith(b"\n"): + parts.append(b"\n") + parts.append(b"endobj\n") + xref_offset = sum(len(p) for p in parts) + max_obj = max(objects.keys()) + parts.append(f"xref\n0 {max_obj + 1}\n".encode()) + parts.append(b"0000000000 65535 f \n") + for i in range(1, max_obj + 1): + off = offsets.get(i, 0) + parts.append(f"{off:010d} 00000 n \n".encode()) + parts.append(f"trailer\n<< /Size {max_obj + 1} /Root 1 0 R >>\n".encode()) + parts.append(f"startxref\n{xref_offset}\n%%EOF\n".encode()) + return b"".join(parts) + + +def _pdf_single_no_annot( + data: bytes = b"Hello Attachment\n", name: str = "hello.txt" +) -> bytes: + ef_len = len(data) + return _build_pdf( + { + 1: b"<< /Type /Catalog /Pages 2 0 R /Names << /EmbeddedFiles << /Names [ (" + + name.encode() + + b") 4 0 R ] >> >> >>", + 2: b"<< /Type /Pages /Kids [ 3 0 R ] /Count 1 >>", + 3: b"<< /Type /Page /Parent 2 0 R /MediaBox [ 0 0 612 792 ] >>", + 4: b"<< /Type /Filespec /F (" + + name.encode() + + b") /UF (" + + name.encode() + + b") /EF << /F 5 0 R /UF 5 0 R >> >>", + 5: b"<< /Type /EmbeddedFile /Subtype /text#2fplain /Length " + + str(ef_len).encode() + + b" /Params << /Size " + + str(ef_len).encode() + + b" >> >>\nstream\n" + + data + + b"\nendstream", + } + ) + + +def _pdf_one_file_two_annots(data: bytes = b"Hello\n", name: str = "doc.txt") -> bytes: + ef_len = len(data) + return _build_pdf( + { + 1: b"<< /Type /Catalog /Pages 2 0 R /Names << /EmbeddedFiles << /Names [ (" + + name.encode() + + b") 4 0 R ] >> >> >>", + 2: b"<< /Type /Pages /Kids [ 3 0 R 6 0 R 7 0 R ] /Count 3 >>", + 3: b"<< /Type /Page /Parent 2 0 R /MediaBox [ 0 0 612 792 ] /Annots [ 8 0 R ] >>", + 4: b"<< /Type /Filespec /F (" + + name.encode() + + b") /UF (" + + name.encode() + + b") /EF << /F 5 0 R /UF 5 0 R >> >>", + 5: b"<< /Type /EmbeddedFile /Length " + + str(ef_len).encode() + + b" /Params << /Size " + + str(ef_len).encode() + + b" >> >>\nstream\n" + + data + + b"\nendstream", + 6: b"<< /Type /Page /Parent 2 0 R /MediaBox [ 0 0 612 792 ] >>", + 7: b"<< /Type /Page /Parent 2 0 R /MediaBox [ 0 0 612 792 ] /Annots [ 9 0 R ] >>", + 8: b"<< /Type /Annot /Subtype /FileAttachment /Rect [ 100 100 120 120 ] /FS 4 0 R /Name /Paperclip >>", + 9: b"<< /Type /Annot /Subtype /FileAttachment /Rect [ 200 200 220 220 ] /FS 4 0 R /Name /Paperclip >>", + } + ) + + +def _pdf_same_name_different_bytes(name: str = "dup.txt") -> bytes: + data_a = b"hello" + data_b = b"world" + return _build_pdf( + { + 1: b"<< /Type /Catalog /Pages 2 0 R /Names << /EmbeddedFiles << /Names [ (" + + name.encode() + + b") 4 0 R (" + + name.encode() + + b") 6 0 R ] >> >> >>", + 2: b"<< /Type /Pages /Kids [ 3 0 R ] /Count 1 >>", + 3: b"<< /Type /Page /Parent 2 0 R /MediaBox [ 0 0 612 792 ] >>", + 4: b"<< /Type /Filespec /F (" + name.encode() + b") /EF << /F 5 0 R >> >>", + 5: b"<< /Length " + + str(len(data_a)).encode() + + b" /Params << /Size " + + str(len(data_a)).encode() + + b" >> >>\nstream\n" + + data_a + + b"\nendstream", + 6: b"<< /Type /Filespec /F (" + name.encode() + b") /EF << /F 7 0 R >> >>", + 7: b"<< /Length " + + str(len(data_b)).encode() + + b" /Params << /Size " + + str(len(data_b)).encode() + + b" >> >>\nstream\n" + + data_b + + b"\nendstream", + } + ) + + +def test_single_attachment_no_annot(): + parser = DoclingPdfParser(loglevel="fatal") + pdf_bytes = _pdf_single_no_annot() + doc = parser.load(path_or_stream=BytesIO(pdf_bytes)) + + atts = doc.get_attachments() + assert len(atts) == 1 + att = atts[0] + assert att.name == "hello.txt" + assert att.mime_type == "text/plain" + assert att.size == len(b"Hello Attachment\n") + assert att.annotations == [] + + # streaming requires max_size + with pytest.raises(TypeError): + doc.get_attachment_data(0) # type: ignore[call-arg] + + data = doc.get_attachment_data(0, max_size=10_000) + assert data == b"Hello Attachment\n" + + # BinaryIO path + bio = doc.get_attachment_stream(0, max_size=10_000) + try: + assert bio.read() == b"Hello Attachment\n" + finally: + bio.close() + + +def test_one_file_two_annots(): + parser = DoclingPdfParser(loglevel="fatal") + pdf_bytes = _pdf_one_file_two_annots() + doc = parser.load(path_or_stream=BytesIO(pdf_bytes)) + + atts = doc.get_attachments() + assert len(atts) == 1, f"ID dedup should collapse to 1, got {len(atts)}" + att = atts[0] + assert att.name == "doc.txt" + assert len(att.annotations) == 2 + # 1-based page numbers (PDF convention): annots on page 1 and 3 + pages = sorted(a.page_no for a in att.annotations) + assert pages == [1, 3] + # exact Rect values from builder, sorted by page_no + annots_sorted = sorted(att.annotations, key=lambda a: a.page_no) + assert annots_sorted[0].page_no == 1 + assert annots_sorted[1].page_no == 3 + b0 = annots_sorted[0].bbox + b1 = annots_sorted[1].bbox + assert (b0.r_x0, b0.r_y0, b0.r_x1, b0.r_y1, b0.r_x2, b0.r_y2, b0.r_x3, b0.r_y3) == ( + 100.0, + 100.0, + 120.0, + 100.0, + 120.0, + 120.0, + 100.0, + 120.0, + ) + assert (b1.r_x0, b1.r_y0, b1.r_x1, b1.r_y1, b1.r_x2, b1.r_y2, b1.r_x3, b1.r_y3) == ( + 200.0, + 200.0, + 220.0, + 200.0, + 220.0, + 220.0, + 200.0, + 220.0, + ) + assert tuple(annots_sorted[0].bbox.to_polygon()[0]) == (100.0, 100.0) + assert tuple(annots_sorted[1].bbox.to_polygon()[0]) == (200.0, 200.0) + assert [tuple(p) for p in annots_sorted[0].bbox.to_polygon()] == [ + (100.0, 100.0), + (120.0, 100.0), + (120.0, 120.0), + (100.0, 120.0), + ] + assert [tuple(p) for p in annots_sorted[1].bbox.to_polygon()] == [ + (200.0, 200.0), + (220.0, 200.0), + (220.0, 220.0), + (200.0, 220.0), + ] + + data = doc.get_attachment_data(0, max_size=10_000) + assert data == b"Hello\n" + + +def test_same_name_different_bytes(): + parser = DoclingPdfParser(loglevel="fatal") + pdf_bytes = _pdf_same_name_different_bytes() + doc = parser.load(path_or_stream=BytesIO(pdf_bytes)) + + atts = doc.get_attachments() + assert len(atts) == 2, ( + f"same name different bytes must stay separate, got {len(atts)}" + ) + names = [a.name for a in atts] + assert names == ["dup.txt", "dup.txt"] + + # Data differs — proves no hash collapse in parse + d0 = doc.get_attachment_data(0, max_size=10_000) + d1 = doc.get_attachment_data(1, max_size=10_000) + assert d0 == b"hello" + assert d1 == b"world" + assert d0 != d1 + + +def test_large_requires_limit(): + parser = DoclingPdfParser(loglevel="fatal") + large = b"x" * (6 * 1024 * 1024) # 6 MB + pdf_bytes = _pdf_single_no_annot(data=large, name="big.bin") + doc = parser.load(path_or_stream=BytesIO(pdf_bytes)) + + atts = doc.get_attachments() + assert len(atts) == 1 + assert atts[0].size == len(large) + + with pytest.raises(Exception, match="max_size"): + doc.get_attachment_data(0, max_size=1024) + + # Should succeed with large enough limit + data = doc.get_attachment_data(0, max_size=10_000_000) + assert len(data) == len(large) + + # stream spill path (>8MB threshold) — use 9 MB to trigger tempfile + huge = b"y" * (9 * 1024 * 1024) + pdf_bytes2 = _pdf_single_no_annot(data=huge, name="huge.bin") + doc2 = parser.load(path_or_stream=BytesIO(pdf_bytes2)) + bio = doc2.get_attachment_stream(0, max_size=10_000_000) + try: + # Should be a file, not BytesIO, for >8MB + assert hasattr(bio, "read") + assert len(bio.read()) == len(huge) + finally: + bio.close() + # cleanup tempfile if created + if hasattr(bio, "name") and isinstance(bio.name, str): + try: + os.unlink(bio.name) + except Exception: + pass + + +def test_empty_pdf_no_attachments(): + parser = DoclingPdfParser(loglevel="fatal") + pdf_bytes = _build_pdf( + { + 1: b"<< /Type /Catalog /Pages 2 0 R >>", + 2: b"<< /Type /Pages /Kids [ 3 0 R ] /Count 1 >>", + 3: b"<< /Type /Page /Parent 2 0 R /MediaBox [ 0 0 612 792 ] >>", + } + ) + doc = parser.load(path_or_stream=BytesIO(pdf_bytes)) + assert doc.get_attachments() == []