diff --git a/napi/src/lib.rs b/napi/src/lib.rs index c4d968d..652469a 100644 --- a/napi/src/lib.rs +++ b/napi/src/lib.rs @@ -89,6 +89,11 @@ pub struct TextItem { pub item_type: ItemType, /// URL for link items, `None` for other types. pub link_url: Option, + /// Marked Content ID from the content stream's BDC/BMC operator, `None` + /// when the text is not part of marked content. Join with the + /// `page`/`mcid` pairs from [`extractStructureElements`] to attach + /// structure-tree roles (headings, paragraphs, …) in tagged PDFs. + pub mcid: Option, } /// A page's regions for text extraction: (page_index_0based, bboxes). @@ -306,12 +311,61 @@ pub fn extract_text_with_positions( is_strikeout: item.is_strikeout, item_type, link_url, + mcid: item.mcid, } }) .collect()) }) } +/// One structure-tree element reference from a tagged PDF. +#[napi(object)] +pub struct StructureElementJs { + /// 1-indexed page number (matches `TextItem.page`). + pub page: u32, + /// Marked Content ID from the page's content stream (matches + /// `TextItem.mcid`). + pub mcid: i64, + /// Standard structure type name ("H1".."H6", "P", "Table", "TD", …). + /// Custom tags are resolved through the document's role map; tags with + /// no standard mapping are returned verbatim. + pub role: String, +} + +/// Extract structure-tree element references from a tagged PDF. +/// +/// Parses the document's structure tree (when present) and returns one +/// entry per marked-content reference, resolved to its 1-indexed page, +/// MCID, and structure type name. Returns an empty array when the PDF is +/// not tagged. +/// +/// Join `(page, mcid)` against the `page`/`mcid` fields from +/// [`extractTextWithPositions`] to attach heading levels (H1..H6) and other +/// semantic roles to extracted text. +/// +/// Pass 1-indexed page numbers (matching `TextItem.page`) to restrict +/// output; omit `pages` for the whole document. Entries are sorted by +/// `(page, mcid)`. +#[napi] +pub fn extract_structure_elements( + buffer: Buffer, + pages: Option>, +) -> Result> { + let bytes: Vec = buffer.to_vec(); + catch_panic("extract_structure_elements", move || { + let elements = pdf_inspector::extract_structure_elements_mem(&bytes, pages.as_deref()) + .map_err(|e| to_napi_err(e, "extract_structure_elements"))?; + Ok(elements + .into_iter() + .map(|e| StructureElementJs { + page: e.page, + mcid: e.mcid, + role: e.role, + }) + .collect()) + }) +} + /// Extract text within bounding-box regions from a PDF. /// /// For hybrid OCR: layout model detects regions in rendered images, diff --git a/napi/test.mjs b/napi/test.mjs index a020274..15585a3 100644 --- a/napi/test.mjs +++ b/napi/test.mjs @@ -8,6 +8,7 @@ import { classifyPdfAsync, extractText, extractTextWithPositions, + extractStructureElements, extractTextInRegions, detectVectorGridInRegion, extractPagesMarkdown, @@ -15,6 +16,7 @@ import { } from './index.js'; const fixture = readFileSync('../tests/fixtures/thermo-freon12.pdf'); +const taggedFixture = readFileSync('../tests/fixtures/firecrawl_docs_tagged.pdf'); // --- processPdf --- console.log('Testing processPdf...'); @@ -82,6 +84,46 @@ assert.ok(page1Items.length > 0); assert.ok(page1Items.every(i => i.page === 1)); console.log(' extractTextWithPositions with pages: OK'); +// mcid: undefined on untagged PDFs, numeric on tagged marked content +assert.ok(items.every(i => i.mcid === undefined || typeof i.mcid === 'number')); +const taggedItems = extractTextWithPositions(taggedFixture); +assert.ok( + taggedItems.some(i => typeof i.mcid === 'number'), + 'tagged PDF text items should carry Marked Content IDs', +); +console.log(' extractTextWithPositions mcid: OK'); + +// --- extractStructureElements --- +console.log('Testing extractStructureElements...'); +const structureElements = extractStructureElements(taggedFixture); +assert.ok(structureElements.length > 0); +assert.ok(structureElements.every(e => typeof e.page === 'number')); +assert.ok(structureElements.every(e => typeof e.mcid === 'number')); +assert.ok(structureElements.every(e => typeof e.role === 'string' && e.role.length > 0)); +assert.ok( + structureElements.some(e => e.role === 'H1'), + 'tagged fixture should surface H1 heading roles', +); + +// (page, mcid) joins against extractTextWithPositions to recover heading text +const h1Refs = new Set( + structureElements.filter(e => e.role === 'H1').map(e => `${e.page}:${e.mcid}`), +); +const h1Text = taggedItems + .filter(i => typeof i.mcid === 'number' && h1Refs.has(`${i.page}:${i.mcid}`)) + .map(i => i.text) + .join(''); +assert.ok(h1Text.trim().length > 0, 'H1 join should recover heading text'); + +// pages filter is 1-indexed, matching TextItem.page +const page1Elements = extractStructureElements(taggedFixture, [1]); +assert.ok(page1Elements.length > 0); +assert.ok(page1Elements.every(e => e.page === 1)); + +// untagged PDFs yield an empty array +assert.deepEqual(extractStructureElements(fixture), []); +console.log(' extractStructureElements: OK'); + // --- extractTextInRegions --- console.log('Testing extractTextInRegions...'); const regionResults = extractTextInRegions(fixture, [ diff --git a/pdf_inspector.pyi b/pdf_inspector.pyi index aa3689e..f2ea3d5 100644 --- a/pdf_inspector.pyi +++ b/pdf_inspector.pyi @@ -51,6 +51,20 @@ class TextItem: is_underline: bool is_strikeout: bool item_type: str + mcid: Optional[int] + """Marked Content ID from the content stream's BDC/BMC operator, None when + the text is not part of marked content. Join with the (page, mcid) pairs + from extract_structure_elements to attach structure-tree roles in tagged + PDFs.""" + +class StructureElement: + """One structure-tree element reference from a tagged PDF.""" + page: int + """1-indexed page number (matches TextItem.page).""" + mcid: int + """Marked Content ID from the page's content stream (matches TextItem.mcid).""" + role: str + """Standard structure type name ("H1".."H6", "P", "Table", "TD", ...).""" class RegionText: """Extracted text for a single region.""" @@ -132,6 +146,27 @@ def extract_text_with_positions_bytes(data: bytes, pages: Optional[list[int]] = """Extract text with position information from bytes.""" ... +def extract_structure_elements(path: str, pages: Optional[list[int]] = None) -> list[StructureElement]: + """Extract structure-tree element references from a tagged PDF file. + + Returns one entry per marked-content reference, resolved to its 1-indexed + page, MCID, and structure type name ("H1".."H6", "P", "Table", ...), sorted + by (page, mcid). Returns an empty list when the PDF is not tagged. + + Args: + path: Path to the PDF file. + pages: Optional list of 1-indexed pages (matching ``TextItem.page``). + When ``None`` (default), the whole document is returned. + """ + ... + +def extract_structure_elements_bytes(data: bytes, pages: Optional[list[int]] = None) -> list[StructureElement]: + """Extract structure-tree element references from tagged PDF bytes. + + See :func:`extract_structure_elements` for details. + """ + ... + def extract_text_in_regions( path: str, page_regions: list[tuple[int, list[list[float]]]], diff --git a/src/lib.rs b/src/lib.rs index 0c9f504..2801294 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -657,6 +657,80 @@ pub fn extract_pages_markdown>( extract_pages_markdown_mem(&buffer, pages) } +// ========================================================================= +// Structure-tree element extraction (tagged PDFs) +// ========================================================================= + +/// One structure-tree element reference from a tagged PDF, resolved to a +/// page and Marked Content ID. +/// +/// Join `(page, mcid)` against [`TextItem::page`] / [`TextItem::mcid`] from +/// [`extract_text_with_positions`] to attach semantic roles (heading levels, +/// paragraphs, table cells, …) to extracted text. +#[derive(Debug, Clone)] +pub struct StructureElement { + /// 1-indexed page number (matches [`TextItem::page`]). + pub page: u32, + /// Marked Content ID from the page's content stream (matches + /// [`TextItem::mcid`]). + pub mcid: i64, + /// Standard structure type name ("H1".."H6", "P", "Table", "TD", …). + /// Custom tags are resolved through the document's `/RoleMap`; tags + /// with no standard mapping are returned verbatim. + pub role: String, +} + +/// Extract structure-tree element references from a tagged PDF in memory. +/// +/// Parses `/StructTreeRoot` (when present) and returns one entry per +/// marked-content reference, resolved to its 1-indexed page, MCID, and +/// structure type name. Returns an empty list when the PDF is not tagged. +/// +/// Pass `Some(&[...])` with 1-indexed page numbers (matching +/// [`TextItem::page`]) to restrict output to those pages; pass `None` for +/// the whole document. Entries are sorted by `(page, mcid)`. +pub fn extract_structure_elements_mem( + buffer: &[u8], + pages: Option<&[u32]>, +) -> Result, PdfError> { + validate_pdf_bytes(buffer)?; + let (doc, _page_count) = load_document_from_mem(buffer)?; + let Some(tree) = structure_tree::StructTree::from_doc(&doc) else { + return Ok(Vec::new()); + }; + let page_ids = doc.get_pages(); + let roles = tree.mcid_to_roles(&page_ids); + + let page_filter: Option> = pages.map(|p| p.iter().copied().collect()); + let mut elements: Vec = roles + .into_iter() + .filter(|(page, _)| page_filter.as_ref().is_none_or(|f| f.contains(page))) + .flat_map(|(page, mcids)| { + mcids.into_iter().map(move |(mcid, role)| StructureElement { + page, + mcid, + role: role.name().to_string(), + }) + }) + .collect(); + elements.sort_unstable_by_key(|e| (e.page, e.mcid)); + Ok(elements) +} + +/// Path-based wrapper for [`extract_structure_elements_mem`]. +/// +/// Reads the PDF from disk and extracts structure-tree element references. +/// Pass `None` for `pages` to return the whole document, or `Some(&[...])` +/// to restrict to specific 1-indexed pages. +pub fn extract_structure_elements>( + path: P, + pages: Option<&[u32]>, +) -> Result, PdfError> { + validate_pdf_file(&path)?; + let buffer = std::fs::read(path.as_ref())?; + extract_structure_elements_mem(&buffer, pages) +} + // ========================================================================= // Region-based text extraction (for hybrid OCR pipelines) // ========================================================================= diff --git a/src/python.rs b/src/python.rs index 8269263..e0e709d 100644 --- a/src/python.rs +++ b/src/python.rs @@ -271,6 +271,12 @@ pub struct PyTextItem { pub is_strikeout: bool, #[pyo3(get)] pub item_type: String, + /// Marked Content ID from the content stream's BDC/BMC operator, None + /// when the text is not part of marked content. Join with the + /// (page, mcid) pairs from extract_structure_elements to attach + /// structure-tree roles (headings, paragraphs, ...) in tagged PDFs. + #[pyo3(get)] + pub mcid: Option, } #[pymethods] @@ -286,6 +292,32 @@ impl PyTextItem { } } +/// One structure-tree element reference from a tagged PDF. +#[pyclass(name = "StructureElement")] +#[derive(Clone)] +pub struct PyStructureElement { + /// 1-indexed page number (matches TextItem.page). + #[pyo3(get)] + pub page: u32, + /// Marked Content ID from the page's content stream (matches + /// TextItem.mcid). + #[pyo3(get)] + pub mcid: i64, + /// Standard structure type name ("H1".."H6", "P", "Table", "TD", ...). + #[pyo3(get)] + pub role: String, +} + +#[pymethods] +impl PyStructureElement { + fn __repr__(&self) -> String { + format!( + "StructureElement(page={}, mcid={}, role='{}')", + self.page, self.mcid, self.role + ) + } +} + // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- @@ -356,6 +388,18 @@ fn convert_text_items(items: Vec) -> Vec { is_underline: item.is_underline, is_strikeout: item.is_strikeout, item_type: item_type_str(&item.item_type), + mcid: item.mcid, + }) + .collect() +} + +fn convert_structure_elements(elements: Vec) -> Vec { + elements + .into_iter() + .map(|e| PyStructureElement { + page: e.page, + mcid: e.mcid, + role: e.role, }) .collect() } @@ -613,6 +657,48 @@ fn extract_pages_markdown_bytes( Ok(to_py_pages_result(result)) } +/// Extract structure-tree element references from a tagged PDF file. +/// +/// Parses the document's structure tree (when present) and returns one +/// entry per marked-content reference, resolved to its 1-indexed page, +/// MCID, and structure type name ("H1".."H6", "P", "Table", ...). Returns +/// an empty list when the PDF is not tagged. +/// +/// Join (page, mcid) against the page/mcid attributes from +/// [`extract_text_with_positions`] to attach heading levels and other +/// semantic roles to extracted text. +/// +/// Args: +/// path: Path to the PDF file. +/// pages: Optional list of 1-indexed pages (matching TextItem.page). +/// When None (default), the whole document is returned. +/// +/// Returns: +/// List of StructureElement sorted by (page, mcid). +#[pyfunction] +#[pyo3(signature = (path, pages=None))] +fn extract_structure_elements( + path: &str, + pages: Option>, +) -> PyResult> { + let elements = crate::extract_structure_elements(path, pages.as_deref()).map_err(to_py_err)?; + Ok(convert_structure_elements(elements)) +} + +/// Extract structure-tree element references from tagged PDF bytes. +/// +/// See [`extract_structure_elements`] for details. +#[pyfunction] +#[pyo3(signature = (data, pages=None))] +fn extract_structure_elements_bytes( + data: &[u8], + pages: Option>, +) -> PyResult> { + let elements = + crate::extract_structure_elements_mem(data, pages.as_deref()).map_err(to_py_err)?; + Ok(convert_structure_elements(elements)) +} + /// Python module definition. #[pymodule] fn pdf_inspector(m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -620,6 +706,7 @@ fn pdf_inspector(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; @@ -634,6 +721,8 @@ fn pdf_inspector(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(extract_text_bytes, m)?)?; m.add_function(wrap_pyfunction!(extract_text_with_positions, m)?)?; m.add_function(wrap_pyfunction!(extract_text_with_positions_bytes, m)?)?; + m.add_function(wrap_pyfunction!(extract_structure_elements, m)?)?; + m.add_function(wrap_pyfunction!(extract_structure_elements_bytes, m)?)?; m.add_function(wrap_pyfunction!(extract_text_in_regions, m)?)?; m.add_function(wrap_pyfunction!(extract_text_in_regions_bytes, m)?)?; m.add_function(wrap_pyfunction!(extract_pages_markdown, m)?)?; diff --git a/src/structure_tree.rs b/src/structure_tree.rs index 44dfeaa..cfa91fd 100644 --- a/src/structure_tree.rs +++ b/src/structure_tree.rs @@ -122,6 +122,65 @@ impl StructRole { ) } + /// The standard structure type name for this role ("H1", "P", "Table", …). + /// + /// Inverse of [`StructRole::from_name`]: for [`StructRole::Other`] the + /// custom tag name is returned verbatim. + pub fn name(&self) -> &str { + match self { + Self::Document => "Document", + Self::Part => "Part", + Self::Art => "Art", + Self::Sect => "Sect", + Self::Div => "Div", + Self::BlockQuote => "BlockQuote", + Self::Caption => "Caption", + Self::TOC => "TOC", + Self::TOCI => "TOCI", + Self::Index => "Index", + Self::NonStruct => "NonStruct", + Self::Private => "Private", + Self::H => "H", + Self::H1 => "H1", + Self::H2 => "H2", + Self::H3 => "H3", + Self::H4 => "H4", + Self::H5 => "H5", + Self::H6 => "H6", + Self::P => "P", + Self::L => "L", + Self::LI => "LI", + Self::Lbl => "Lbl", + Self::LBody => "LBody", + Self::Table => "Table", + Self::TR => "TR", + Self::TH => "TH", + Self::TD => "TD", + Self::THead => "THead", + Self::TBody => "TBody", + Self::TFoot => "TFoot", + Self::Span => "Span", + Self::Quote => "Quote", + Self::Note => "Note", + Self::Reference => "Reference", + Self::BibEntry => "BibEntry", + Self::Code => "Code", + Self::Link => "Link", + Self::Annot => "Annot", + Self::Figure => "Figure", + Self::Formula => "Formula", + Self::Form => "Form", + Self::Ruby => "Ruby", + Self::RB => "RB", + Self::RT => "RT", + Self::RP => "RP", + Self::Warichu => "Warichu", + Self::WT => "WT", + Self::WP => "WP", + Self::Other(name) => name, + } + } + fn from_name(name: &str) -> Self { match name { "Document" => Self::Document, @@ -1233,6 +1292,66 @@ mod tests { ); } + #[test] + fn test_struct_role_name_roundtrip() { + // `name()` is the inverse of `from_name` for every standard type. + for name in [ + "Document", + "Part", + "Art", + "Sect", + "Div", + "BlockQuote", + "Caption", + "TOC", + "TOCI", + "Index", + "NonStruct", + "Private", + "H", + "H1", + "H2", + "H3", + "H4", + "H5", + "H6", + "P", + "L", + "LI", + "Lbl", + "LBody", + "Table", + "TR", + "TH", + "TD", + "THead", + "TBody", + "TFoot", + "Span", + "Quote", + "Note", + "Reference", + "BibEntry", + "Code", + "Link", + "Annot", + "Figure", + "Formula", + "Form", + "Ruby", + "RB", + "RT", + "RP", + "Warichu", + "WT", + "WP", + ] { + assert_eq!(StructRole::from_name(name).name(), name); + } + // Custom tags pass through verbatim. + assert_eq!(StructRole::from_name("CustomTag").name(), "CustomTag"); + } + #[test] fn test_struct_role_with_role_map() { let mut role_map = HashMap::new(); diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs index 0506291..52cfce7 100644 --- a/tests/integration_tests.rs +++ b/tests/integration_tests.rs @@ -1429,6 +1429,77 @@ fn test_firecrawl_tagged_pdf_struct_tree() { assert_eq!(fence_count % 2, 0, "Code fences should be balanced"); } +#[test] +fn test_tagged_pdf_text_items_carry_mcid() { + let buf = std::fs::read("tests/fixtures/firecrawl_docs_tagged.pdf").unwrap(); + let items = pdf_inspector::extractor::extract_text_with_positions_mem(&buf).unwrap(); + assert!( + items.iter().any(|i| i.mcid.is_some()), + "Tagged PDF text items should carry Marked Content IDs" + ); +} + +#[test] +fn test_extract_structure_elements_tagged_pdf() { + let buf = std::fs::read("tests/fixtures/firecrawl_docs_tagged.pdf").unwrap(); + let elements = pdf_inspector::extract_structure_elements_mem(&buf, None).unwrap(); + assert!(!elements.is_empty(), "Tagged PDF should yield elements"); + assert!( + elements.iter().any(|e| e.role == "H1"), + "Should surface H1 heading roles" + ); + assert!( + elements.iter().all(|e| !e.role.is_empty()), + "Every element should carry a role name" + ); + + // Sorted by (page, mcid) for deterministic output + assert!( + elements + .windows(2) + .all(|w| (w[0].page, w[0].mcid) <= (w[1].page, w[1].mcid)), + "Elements should be sorted by (page, mcid)" + ); + + // The advertised join: (page, mcid) pairs must line up with the + // mcid-carrying TextItems from positioned extraction, and joining the + // H1 entries must recover non-empty heading text. + let items = pdf_inspector::extractor::extract_text_with_positions_mem(&buf).unwrap(); + let h1_refs: std::collections::HashSet<(u32, i64)> = elements + .iter() + .filter(|e| e.role == "H1") + .map(|e| (e.page, e.mcid)) + .collect(); + let h1_text: String = items + .iter() + .filter(|i| i.mcid.is_some_and(|mcid| h1_refs.contains(&(i.page, mcid)))) + .map(|i| i.text.as_str()) + .collect(); + assert!( + !h1_text.trim().is_empty(), + "Joining H1 structure elements to text items should recover heading text" + ); + + // Page filter is 1-indexed (matching TextItem.page) and equals the + // corresponding subset of the full document result. + let page1 = pdf_inspector::extract_structure_elements_mem(&buf, Some(&[1])).unwrap(); + assert!(!page1.is_empty(), "Page 1 should have elements"); + assert!(page1.iter().all(|e| e.page == 1)); + let full_page1_count = elements.iter().filter(|e| e.page == 1).count(); + assert_eq!(page1.len(), full_page1_count); +} + +#[test] +fn test_extract_structure_elements_untagged_pdf_empty() { + let buf = std::fs::read("tests/fixtures/thermo-freon12.pdf").unwrap(); + let elements = pdf_inspector::extract_structure_elements_mem(&buf, None).unwrap(); + assert!( + elements.is_empty(), + "Untagged PDF should yield no structure elements, got {:?}", + elements + ); +} + #[test] fn test_identity_h_no_tounicode_suppresses_garbage() { // shinagawa_identity_h.pdf uses YuGothic with Identity-H encoding and no diff --git a/tests/test_python.py b/tests/test_python.py index a17e1d7..ccf9752 100644 --- a/tests/test_python.py +++ b/tests/test_python.py @@ -203,6 +203,79 @@ class TestExtractTextWithPositions: assert len(items) > 0 assert all(item.page == 1 for item in items) + def test_mcid(self): + # Untagged fixture: mcid is None or int, never anything else + items = pdf_inspector.extract_text_with_positions( + fixture_path("thermo-freon12.pdf") + ) + assert all(item.mcid is None or isinstance(item.mcid, int) for item in items) + # Tagged fixture: marked content carries MCIDs + tagged = pdf_inspector.extract_text_with_positions( + fixture_path("firecrawl_docs_tagged.pdf") + ) + assert any(item.mcid is not None for item in tagged) + + +# --------------------------------------------------------------------------- +# extract_structure_elements / extract_structure_elements_bytes +# --------------------------------------------------------------------------- + + +class TestExtractStructureElements: + def test_tagged_file(self): + elements = pdf_inspector.extract_structure_elements( + fixture_path("firecrawl_docs_tagged.pdf") + ) + assert len(elements) > 0 + assert all(isinstance(e.page, int) for e in elements) + assert all(isinstance(e.mcid, int) for e in elements) + assert all(isinstance(e.role, str) and len(e.role) > 0 for e in elements) + assert any(e.role == "H1" for e in elements) + + def test_join_with_text_items(self): + # (page, mcid) joins against extract_text_with_positions to recover + # heading text + path = fixture_path("firecrawl_docs_tagged.pdf") + elements = pdf_inspector.extract_structure_elements(path) + items = pdf_inspector.extract_text_with_positions(path) + h1_refs = {(e.page, e.mcid) for e in elements if e.role == "H1"} + h1_text = "".join( + item.text + for item in items + if item.mcid is not None and (item.page, item.mcid) in h1_refs + ) + assert len(h1_text.strip()) > 0 + + def test_with_pages(self): + # pages filter is 1-indexed, matching TextItem.page + elements = pdf_inspector.extract_structure_elements( + fixture_path("firecrawl_docs_tagged.pdf"), pages=[1] + ) + assert len(elements) > 0 + assert all(e.page == 1 for e in elements) + + def test_bytes(self): + data = fixture_bytes("firecrawl_docs_tagged.pdf") + elements = pdf_inspector.extract_structure_elements_bytes(data) + assert len(elements) > 0 + assert any(e.role == "H1" for e in elements) + + def test_untagged_returns_empty(self): + elements = pdf_inspector.extract_structure_elements( + fixture_path("thermo-freon12.pdf") + ) + assert elements == [] + + def test_repr(self): + elements = pdf_inspector.extract_structure_elements( + fixture_path("firecrawl_docs_tagged.pdf") + ) + assert "StructureElement" in repr(elements[0]) + + def test_not_a_pdf(self): + with pytest.raises(ValueError): + pdf_inspector.extract_structure_elements_bytes(b"not a pdf") + # --------------------------------------------------------------------------- # extract_text_in_regions / extract_text_in_regions_bytes