feat: add Python bindings via PyO3
Expose the pdf-inspector Rust library as a Python package using PyO3 + maturin. Python users can now `pip install` and use `import pdf_inspector` for PDF classification, text extraction, and markdown conversion with native Rust speed. Adds: - src/python.rs: PyO3 bindings (process_pdf, detect_pdf, extract_text, etc.) - pyproject.toml: maturin build configuration - pdf_inspector.pyi: type stubs for IDE support - tests/test_python.py: 21 pytest tests covering all Python API functions - examples/basic_usage.py: example script demonstrating all features - Updated README with Python quick start and API reference Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Jacques Dumora
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
7a0e074fa2
commit
0bf5463a1e
@@ -8,7 +8,14 @@ description = "Fast PDF inspection, classification, and text extraction with sma
|
||||
license = "MIT"
|
||||
repository = "https://github.com/firecrawl/pdf-inspector"
|
||||
|
||||
[lib]
|
||||
name = "pdf_inspector"
|
||||
crate-type = ["lib", "cdylib"]
|
||||
|
||||
[dependencies]
|
||||
# Python bindings
|
||||
pyo3 = { version = "0.22", features = ["extension-module"], optional = true }
|
||||
|
||||
# PDF parsing
|
||||
lopdf = { git = "https://github.com/firecrawl/lopdf", branch = "firecrawl/zlib-checksum-encrypted", features = ["rayon"] }
|
||||
|
||||
@@ -34,6 +41,7 @@ tempfile = "3.3"
|
||||
|
||||
[features]
|
||||
default = []
|
||||
python = ["pyo3"]
|
||||
|
||||
[[bin]]
|
||||
name = "pdf2md"
|
||||
|
||||
@@ -15,10 +15,70 @@ Built by [Firecrawl](https://firecrawl.dev) to handle text-based PDFs locally in
|
||||
- **Encoding issue detection** — Automatically flags broken font encodings (garbled text, replacement characters) so callers can fall back to OCR.
|
||||
- **Single document load** — The document is parsed once and shared between detection and extraction, avoiding redundant I/O.
|
||||
- **Lightweight** — Pure Rust, no ML models, no external services. Single dependency on `lopdf` for PDF parsing.
|
||||
- **Python bindings** — Use from Python via PyO3. Install with `pip install pdf-inspector` or build from source with `maturin`.
|
||||
|
||||
## Quick start
|
||||
|
||||
### As a library
|
||||
### Python
|
||||
|
||||
Install from source (requires Rust toolchain):
|
||||
|
||||
```bash
|
||||
pip install maturin
|
||||
maturin develop --release
|
||||
```
|
||||
|
||||
Use it:
|
||||
|
||||
```python
|
||||
import pdf_inspector
|
||||
|
||||
# Full processing: detect + extract + convert to Markdown
|
||||
result = pdf_inspector.process_pdf("document.pdf")
|
||||
print(result.pdf_type) # "text_based", "scanned", "image_based", "mixed"
|
||||
print(result.confidence) # 0.0 - 1.0
|
||||
print(result.page_count) # number of pages
|
||||
print(result.markdown) # Markdown string or None
|
||||
|
||||
# Process specific pages only
|
||||
result = pdf_inspector.process_pdf("document.pdf", pages=[1, 3, 5])
|
||||
|
||||
# Process from bytes (no filesystem needed)
|
||||
with open("document.pdf", "rb") as f:
|
||||
result = pdf_inspector.process_pdf_bytes(f.read())
|
||||
|
||||
# Fast detection only (no text extraction)
|
||||
result = pdf_inspector.detect_pdf("document.pdf")
|
||||
if result.pdf_type == "text_based":
|
||||
print("Can extract locally!")
|
||||
else:
|
||||
print(f"Pages needing OCR: {result.pages_needing_ocr}")
|
||||
|
||||
# Plain text extraction
|
||||
text = pdf_inspector.extract_text("document.pdf")
|
||||
|
||||
# Positioned text items with font info
|
||||
items = pdf_inspector.extract_text_with_positions("document.pdf")
|
||||
for item in items[:5]:
|
||||
print(f"'{item.text}' at ({item.x:.0f}, {item.y:.0f}) size={item.font_size}")
|
||||
```
|
||||
|
||||
#### Python API reference
|
||||
|
||||
| Function | Description |
|
||||
|---|---|
|
||||
| `process_pdf(path, pages=None)` | Full processing (detect + extract + markdown) |
|
||||
| `process_pdf_bytes(data, pages=None)` | Full processing from bytes |
|
||||
| `detect_pdf(path)` | Fast detection only |
|
||||
| `detect_pdf_bytes(data)` | Fast detection from bytes |
|
||||
| `extract_text(path)` | Plain text extraction |
|
||||
| `extract_text_with_positions(path, pages=None)` | Text with X/Y coords and font info |
|
||||
|
||||
**`PdfResult` fields:** `pdf_type`, `markdown`, `page_count`, `processing_time_ms`, `pages_needing_ocr`, `title`, `confidence`, `is_complex_layout`, `pages_with_tables`, `pages_with_columns`, `has_encoding_issues`
|
||||
|
||||
**`TextItem` fields:** `text`, `x`, `y`, `width`, `height`, `font`, `font_size`, `page`, `is_bold`, `is_italic`, `item_type`
|
||||
|
||||
### Rust
|
||||
|
||||
Add to your `Cargo.toml`:
|
||||
|
||||
@@ -159,6 +219,7 @@ The document is loaded **once** via `load_document_from_path` / `load_document_f
|
||||
```
|
||||
src/
|
||||
lib.rs — Public API, PdfOptions builder, convenience functions
|
||||
python.rs — PyO3 Python bindings
|
||||
types.rs — Shared types: TextItem, TextLine, PdfRect, ItemType
|
||||
text_utils.rs — Character/text helpers (CJK, RTL, ligatures, bold/italic)
|
||||
process_mode.rs — ProcessMode enum (DetectOnly, Analyze, Full)
|
||||
@@ -189,7 +250,7 @@ This detects 300+ page PDFs in milliseconds. The result includes `pages_needing_
|
||||
| `Sample(n)` | Sample `n` evenly distributed pages (first, last, middle) | Very large PDFs where speed matters more than precision |
|
||||
| `Pages(vec)` | Only scan specific 1-indexed page numbers | When the caller knows which pages to check |
|
||||
|
||||
## API
|
||||
## Rust API
|
||||
|
||||
### Processing modes
|
||||
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
"""Basic usage examples for pdf-inspector Python library."""
|
||||
|
||||
import sys
|
||||
import pdf_inspector
|
||||
|
||||
|
||||
def main():
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python basic_usage.py <path-to-pdf>")
|
||||
sys.exit(1)
|
||||
|
||||
path = sys.argv[1]
|
||||
|
||||
# 1. Full processing: detect + extract + markdown
|
||||
print("=" * 60)
|
||||
print("Full processing")
|
||||
print("=" * 60)
|
||||
result = pdf_inspector.process_pdf(path)
|
||||
print(f"Type: {result.pdf_type}")
|
||||
print(f"Pages: {result.page_count}")
|
||||
print(f"Confidence: {result.confidence:.0%}")
|
||||
print(f"Time: {result.processing_time_ms}ms")
|
||||
print(f"Title: {result.title}")
|
||||
print(f"Complex: {result.is_complex_layout}")
|
||||
print(f"Tables on: {result.pages_with_tables}")
|
||||
print(f"Columns on: {result.pages_with_columns}")
|
||||
print(f"Encoding: {'issues detected' if result.has_encoding_issues else 'ok'}")
|
||||
print(f"OCR needed: {result.pages_needing_ocr or 'none'}")
|
||||
if result.markdown:
|
||||
print(f"\n--- Markdown ({len(result.markdown)} chars) ---")
|
||||
print(result.markdown[:500])
|
||||
if len(result.markdown) > 500:
|
||||
print(f"\n... ({len(result.markdown) - 500} more chars)")
|
||||
|
||||
# 2. Fast detection only
|
||||
print("\n" + "=" * 60)
|
||||
print("Detection only")
|
||||
print("=" * 60)
|
||||
info = pdf_inspector.detect_pdf(path)
|
||||
print(f"Type: {info.pdf_type}")
|
||||
print(f"Confidence: {info.confidence:.0%}")
|
||||
print(f"Time: {info.processing_time_ms}ms")
|
||||
|
||||
# 3. From bytes
|
||||
print("\n" + "=" * 60)
|
||||
print("From bytes")
|
||||
print("=" * 60)
|
||||
with open(path, "rb") as f:
|
||||
data = f.read()
|
||||
result = pdf_inspector.process_pdf_bytes(data)
|
||||
print(f"Type: {result.pdf_type}, Pages: {result.page_count}")
|
||||
|
||||
# 4. Plain text
|
||||
print("\n" + "=" * 60)
|
||||
print("Plain text extraction")
|
||||
print("=" * 60)
|
||||
text = pdf_inspector.extract_text(path)
|
||||
print(text[:300])
|
||||
|
||||
# 5. Positioned items
|
||||
print("\n" + "=" * 60)
|
||||
print("Positioned text items (first 10)")
|
||||
print("=" * 60)
|
||||
items = pdf_inspector.extract_text_with_positions(path, pages=[1])
|
||||
for item in items[:10]:
|
||||
bold = " [B]" if item.is_bold else ""
|
||||
italic = " [I]" if item.is_italic else ""
|
||||
print(
|
||||
f" p{item.page} ({item.x:6.1f}, {item.y:6.1f}) "
|
||||
f"size={item.font_size:5.1f}{bold}{italic} "
|
||||
f"'{item.text}'"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,56 @@
|
||||
"""Type stubs for pdf_inspector."""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
class PdfResult:
|
||||
"""Result of processing a PDF file."""
|
||||
pdf_type: str
|
||||
"""'text_based', 'scanned', 'image_based', or 'mixed'."""
|
||||
markdown: Optional[str]
|
||||
page_count: int
|
||||
processing_time_ms: int
|
||||
pages_needing_ocr: list[int]
|
||||
title: Optional[str]
|
||||
confidence: float
|
||||
is_complex_layout: bool
|
||||
pages_with_tables: list[int]
|
||||
pages_with_columns: list[int]
|
||||
has_encoding_issues: bool
|
||||
|
||||
class TextItem:
|
||||
"""A positioned text item extracted from a PDF."""
|
||||
text: str
|
||||
x: float
|
||||
y: float
|
||||
width: float
|
||||
height: float
|
||||
font: str
|
||||
font_size: float
|
||||
page: int
|
||||
is_bold: bool
|
||||
is_italic: bool
|
||||
item_type: str
|
||||
|
||||
def process_pdf(path: str, pages: Optional[list[int]] = None) -> PdfResult:
|
||||
"""Process a PDF: detect type, extract text, convert to Markdown."""
|
||||
...
|
||||
|
||||
def process_pdf_bytes(data: bytes, pages: Optional[list[int]] = None) -> PdfResult:
|
||||
"""Process a PDF from bytes in memory."""
|
||||
...
|
||||
|
||||
def detect_pdf(path: str) -> PdfResult:
|
||||
"""Fast detection only — no text extraction."""
|
||||
...
|
||||
|
||||
def detect_pdf_bytes(data: bytes) -> PdfResult:
|
||||
"""Fast detection from bytes."""
|
||||
...
|
||||
|
||||
def extract_text(path: str) -> str:
|
||||
"""Extract plain text from a PDF."""
|
||||
...
|
||||
|
||||
def extract_text_with_positions(path: str, pages: Optional[list[int]] = None) -> list[TextItem]:
|
||||
"""Extract text with position information."""
|
||||
...
|
||||
@@ -0,0 +1,21 @@
|
||||
[build-system]
|
||||
requires = ["maturin>=1.0,<2.0"]
|
||||
build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "pdf-inspector"
|
||||
version = "0.1.0"
|
||||
description = "Fast PDF inspection, classification, and text extraction with smart scanned vs text-based detection"
|
||||
license = { text = "MIT" }
|
||||
requires-python = ">=3.8"
|
||||
classifiers = [
|
||||
"Programming Language :: Rust",
|
||||
"Programming Language :: Python :: Implementation :: CPython",
|
||||
"Programming Language :: Python :: 3",
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Operating System :: OS Independent",
|
||||
"Topic :: Text Processing",
|
||||
]
|
||||
|
||||
[tool.maturin]
|
||||
features = ["python"]
|
||||
@@ -22,6 +22,9 @@
|
||||
//! ).unwrap();
|
||||
//! ```
|
||||
|
||||
#[cfg(feature = "python")]
|
||||
pub mod python;
|
||||
|
||||
pub mod adobe_korea1;
|
||||
pub mod detector;
|
||||
pub mod extractor;
|
||||
|
||||
+269
@@ -0,0 +1,269 @@
|
||||
//! PyO3 Python bindings for pdf-inspector.
|
||||
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use std::collections::HashSet;
|
||||
|
||||
use crate::detector::PdfType;
|
||||
use crate::types::ItemType;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Result wrapper
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Result of processing a PDF file.
|
||||
#[pyclass(name = "PdfResult")]
|
||||
#[derive(Clone)]
|
||||
pub struct PyPdfResult {
|
||||
/// The detected PDF type: "text_based", "scanned", "image_based", or "mixed".
|
||||
#[pyo3(get)]
|
||||
pub pdf_type: String,
|
||||
/// Markdown output (None if detect-only or scanned PDF).
|
||||
#[pyo3(get)]
|
||||
pub markdown: Option<String>,
|
||||
/// Total number of pages.
|
||||
#[pyo3(get)]
|
||||
pub page_count: u32,
|
||||
/// Processing time in milliseconds.
|
||||
#[pyo3(get)]
|
||||
pub processing_time_ms: u64,
|
||||
/// 1-indexed page numbers that need OCR.
|
||||
#[pyo3(get)]
|
||||
pub pages_needing_ocr: Vec<u32>,
|
||||
/// Title from PDF metadata.
|
||||
#[pyo3(get)]
|
||||
pub title: Option<String>,
|
||||
/// Detection confidence (0.0-1.0).
|
||||
#[pyo3(get)]
|
||||
pub confidence: f32,
|
||||
/// Whether the layout is complex (tables/columns detected).
|
||||
#[pyo3(get)]
|
||||
pub is_complex_layout: bool,
|
||||
/// Pages with tables detected.
|
||||
#[pyo3(get)]
|
||||
pub pages_with_tables: Vec<u32>,
|
||||
/// Pages with multi-column layout.
|
||||
#[pyo3(get)]
|
||||
pub pages_with_columns: Vec<u32>,
|
||||
/// Whether encoding issues were detected.
|
||||
#[pyo3(get)]
|
||||
pub has_encoding_issues: bool,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyPdfResult {
|
||||
fn __repr__(&self) -> String {
|
||||
format!(
|
||||
"PdfResult(pdf_type='{}', pages={}, confidence={:.2})",
|
||||
self.pdf_type, self.page_count, self.confidence
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn pdf_type_str(t: PdfType) -> String {
|
||||
match t {
|
||||
PdfType::TextBased => "text_based".into(),
|
||||
PdfType::Scanned => "scanned".into(),
|
||||
PdfType::ImageBased => "image_based".into(),
|
||||
PdfType::Mixed => "mixed".into(),
|
||||
}
|
||||
}
|
||||
|
||||
fn to_py_result(r: crate::PdfProcessResult) -> PyPdfResult {
|
||||
PyPdfResult {
|
||||
pdf_type: pdf_type_str(r.pdf_type),
|
||||
markdown: r.markdown,
|
||||
page_count: r.page_count,
|
||||
processing_time_ms: r.processing_time_ms,
|
||||
pages_needing_ocr: r.pages_needing_ocr,
|
||||
title: r.title,
|
||||
confidence: r.confidence,
|
||||
is_complex_layout: r.layout.is_complex,
|
||||
pages_with_tables: r.layout.pages_with_tables,
|
||||
pages_with_columns: r.layout.pages_with_columns,
|
||||
has_encoding_issues: r.has_encoding_issues,
|
||||
}
|
||||
}
|
||||
|
||||
fn to_py_err(e: crate::PdfError) -> PyErr {
|
||||
PyValueError::new_err(e.to_string())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Text item wrapper
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A positioned text item extracted from a PDF.
|
||||
#[pyclass(name = "TextItem")]
|
||||
#[derive(Clone)]
|
||||
pub struct PyTextItem {
|
||||
#[pyo3(get)]
|
||||
pub text: String,
|
||||
#[pyo3(get)]
|
||||
pub x: f32,
|
||||
#[pyo3(get)]
|
||||
pub y: f32,
|
||||
#[pyo3(get)]
|
||||
pub width: f32,
|
||||
#[pyo3(get)]
|
||||
pub height: f32,
|
||||
#[pyo3(get)]
|
||||
pub font: String,
|
||||
#[pyo3(get)]
|
||||
pub font_size: f32,
|
||||
#[pyo3(get)]
|
||||
pub page: u32,
|
||||
#[pyo3(get)]
|
||||
pub is_bold: bool,
|
||||
#[pyo3(get)]
|
||||
pub is_italic: bool,
|
||||
#[pyo3(get)]
|
||||
pub item_type: String,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyTextItem {
|
||||
fn __repr__(&self) -> String {
|
||||
format!(
|
||||
"TextItem(text='{}', page={}, x={:.1}, y={:.1})",
|
||||
self.text.chars().take(40).collect::<String>(),
|
||||
self.page,
|
||||
self.x,
|
||||
self.y,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn item_type_str(t: &ItemType) -> String {
|
||||
match t {
|
||||
ItemType::Text => "text".into(),
|
||||
ItemType::Image => "image".into(),
|
||||
ItemType::Link(url) => format!("link:{url}"),
|
||||
ItemType::FormField => "form_field".into(),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Public Python API
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Process a PDF file: detect type, extract text, and convert to Markdown.
|
||||
///
|
||||
/// Args:
|
||||
/// path: Path to the PDF file.
|
||||
/// pages: Optional list of 1-indexed page numbers to process.
|
||||
///
|
||||
/// Returns:
|
||||
/// PdfResult with markdown, pdf_type, and metadata.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (path, pages=None))]
|
||||
fn process_pdf(path: &str, pages: Option<Vec<u32>>) -> PyResult<PyPdfResult> {
|
||||
let mut opts = crate::PdfOptions::new();
|
||||
if let Some(p) = pages {
|
||||
opts = opts.pages(p);
|
||||
}
|
||||
let result = crate::process_pdf_with_options(path, opts).map_err(to_py_err)?;
|
||||
Ok(to_py_result(result))
|
||||
}
|
||||
|
||||
/// Process a PDF from bytes in memory.
|
||||
///
|
||||
/// Args:
|
||||
/// data: PDF file contents as bytes.
|
||||
/// pages: Optional list of 1-indexed page numbers to process.
|
||||
///
|
||||
/// Returns:
|
||||
/// PdfResult with markdown, pdf_type, and metadata.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (data, pages=None))]
|
||||
fn process_pdf_bytes(data: &[u8], pages: Option<Vec<u32>>) -> PyResult<PyPdfResult> {
|
||||
let mut opts = crate::PdfOptions::new();
|
||||
if let Some(p) = pages {
|
||||
opts = opts.pages(p);
|
||||
}
|
||||
let result = crate::process_pdf_mem_with_options(data, opts).map_err(to_py_err)?;
|
||||
Ok(to_py_result(result))
|
||||
}
|
||||
|
||||
/// Fast detection only — no text extraction or markdown.
|
||||
///
|
||||
/// Args:
|
||||
/// path: Path to the PDF file.
|
||||
///
|
||||
/// Returns:
|
||||
/// PdfResult with pdf_type and metadata (markdown will be None).
|
||||
#[pyfunction]
|
||||
fn detect_pdf(path: &str) -> PyResult<PyPdfResult> {
|
||||
let result = crate::detect_pdf(path).map_err(to_py_err)?;
|
||||
Ok(to_py_result(result))
|
||||
}
|
||||
|
||||
/// Fast detection from bytes — no text extraction or markdown.
|
||||
#[pyfunction]
|
||||
fn detect_pdf_bytes(data: &[u8]) -> PyResult<PyPdfResult> {
|
||||
let result = crate::detect_pdf_mem(data).map_err(to_py_err)?;
|
||||
Ok(to_py_result(result))
|
||||
}
|
||||
|
||||
/// Extract plain text from a PDF file.
|
||||
///
|
||||
/// Args:
|
||||
/// path: Path to the PDF file.
|
||||
///
|
||||
/// Returns:
|
||||
/// Extracted text as a string.
|
||||
#[pyfunction]
|
||||
fn extract_text(path: &str) -> PyResult<String> {
|
||||
crate::extract_text(path).map_err(to_py_err)
|
||||
}
|
||||
|
||||
/// Extract text with position information.
|
||||
///
|
||||
/// Args:
|
||||
/// path: Path to the PDF file.
|
||||
/// pages: Optional list of 1-indexed page numbers.
|
||||
///
|
||||
/// Returns:
|
||||
/// List of TextItem objects with text, position, font info.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (path, pages=None))]
|
||||
fn extract_text_with_positions(path: &str, pages: Option<Vec<u32>>) -> PyResult<Vec<PyTextItem>> {
|
||||
let items = match pages {
|
||||
Some(p) => {
|
||||
let page_set: HashSet<u32> = p.into_iter().collect();
|
||||
crate::extract_text_with_positions_pages(path, Some(&page_set)).map_err(to_py_err)?
|
||||
}
|
||||
None => crate::extract_text_with_positions(path).map_err(to_py_err)?,
|
||||
};
|
||||
|
||||
Ok(items
|
||||
.into_iter()
|
||||
.map(|item| PyTextItem {
|
||||
text: item.text,
|
||||
x: item.x,
|
||||
y: item.y,
|
||||
width: item.width,
|
||||
height: item.height,
|
||||
font: item.font,
|
||||
font_size: item.font_size,
|
||||
page: item.page,
|
||||
is_bold: item.is_bold,
|
||||
is_italic: item.is_italic,
|
||||
item_type: item_type_str(&item.item_type),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Python module definition.
|
||||
#[pymodule]
|
||||
fn pdf_inspector(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<PyPdfResult>()?;
|
||||
m.add_class::<PyTextItem>()?;
|
||||
m.add_function(wrap_pyfunction!(process_pdf, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(process_pdf_bytes, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(detect_pdf, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(detect_pdf_bytes, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(extract_text, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(extract_text_with_positions, m)?)?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Tests for the pdf_inspector Python bindings."""
|
||||
|
||||
import os
|
||||
import pytest
|
||||
import pdf_inspector
|
||||
|
||||
FIXTURES_DIR = os.path.join(os.path.dirname(__file__), "fixtures")
|
||||
|
||||
|
||||
def fixture_path(name: str) -> str:
|
||||
return os.path.join(FIXTURES_DIR, name)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# process_pdf
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestProcessPdf:
|
||||
def test_basic(self):
|
||||
result = pdf_inspector.process_pdf(fixture_path("thermo-freon12.pdf"))
|
||||
assert result.pdf_type == "text_based"
|
||||
assert result.page_count == 3
|
||||
assert result.confidence > 0.0
|
||||
assert result.markdown is not None
|
||||
assert len(result.markdown) > 0
|
||||
|
||||
def test_result_repr(self):
|
||||
result = pdf_inspector.process_pdf(fixture_path("thermo-freon12.pdf"))
|
||||
r = repr(result)
|
||||
assert "PdfResult" in r
|
||||
assert "text_based" in r
|
||||
|
||||
def test_with_pages(self):
|
||||
result = pdf_inspector.process_pdf(
|
||||
fixture_path("thermo-freon12.pdf"), pages=[1]
|
||||
)
|
||||
assert result.page_count == 3 # total pages in doc
|
||||
assert result.markdown is not None
|
||||
|
||||
def test_result_fields(self):
|
||||
result = pdf_inspector.process_pdf(fixture_path("thermo-freon12.pdf"))
|
||||
# All fields should be accessible
|
||||
assert isinstance(result.pdf_type, str)
|
||||
assert isinstance(result.page_count, int)
|
||||
assert isinstance(result.processing_time_ms, int)
|
||||
assert isinstance(result.pages_needing_ocr, list)
|
||||
assert isinstance(result.confidence, float)
|
||||
assert isinstance(result.is_complex_layout, bool)
|
||||
assert isinstance(result.pages_with_tables, list)
|
||||
assert isinstance(result.pages_with_columns, list)
|
||||
assert isinstance(result.has_encoding_issues, bool)
|
||||
# title can be None or str
|
||||
assert result.title is None or isinstance(result.title, str)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# process_pdf_bytes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestProcessPdfBytes:
|
||||
def test_basic(self):
|
||||
with open(fixture_path("thermo-freon12.pdf"), "rb") as f:
|
||||
data = f.read()
|
||||
result = pdf_inspector.process_pdf_bytes(data)
|
||||
assert result.pdf_type == "text_based"
|
||||
assert result.markdown is not None
|
||||
|
||||
def test_with_pages(self):
|
||||
with open(fixture_path("thermo-freon12.pdf"), "rb") as f:
|
||||
data = f.read()
|
||||
result = pdf_inspector.process_pdf_bytes(data, pages=[1, 2])
|
||||
assert result.markdown is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# detect_pdf / detect_pdf_bytes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDetectPdf:
|
||||
def test_detect_file(self):
|
||||
result = pdf_inspector.detect_pdf(fixture_path("thermo-freon12.pdf"))
|
||||
assert result.pdf_type == "text_based"
|
||||
assert result.markdown is None # detect only — no markdown
|
||||
assert result.page_count == 3
|
||||
|
||||
def test_detect_bytes(self):
|
||||
with open(fixture_path("thermo-freon12.pdf"), "rb") as f:
|
||||
data = f.read()
|
||||
result = pdf_inspector.detect_pdf_bytes(data)
|
||||
assert result.pdf_type == "text_based"
|
||||
assert result.markdown is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# extract_text
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestExtractText:
|
||||
def test_basic(self):
|
||||
text = pdf_inspector.extract_text(fixture_path("thermo-freon12.pdf"))
|
||||
assert isinstance(text, str)
|
||||
assert len(text) > 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# extract_text_with_positions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestExtractTextWithPositions:
|
||||
def test_basic(self):
|
||||
items = pdf_inspector.extract_text_with_positions(
|
||||
fixture_path("thermo-freon12.pdf")
|
||||
)
|
||||
assert len(items) > 0
|
||||
item = items[0]
|
||||
assert isinstance(item.text, str)
|
||||
assert isinstance(item.x, float)
|
||||
assert isinstance(item.y, float)
|
||||
assert isinstance(item.width, float)
|
||||
assert isinstance(item.height, float)
|
||||
assert isinstance(item.font, str)
|
||||
assert isinstance(item.font_size, float)
|
||||
assert isinstance(item.page, int)
|
||||
assert isinstance(item.is_bold, bool)
|
||||
assert isinstance(item.is_italic, bool)
|
||||
assert isinstance(item.item_type, str)
|
||||
|
||||
def test_with_pages(self):
|
||||
items = pdf_inspector.extract_text_with_positions(
|
||||
fixture_path("thermo-freon12.pdf"), pages=[1]
|
||||
)
|
||||
assert len(items) > 0
|
||||
assert all(item.page == 1 for item in items)
|
||||
|
||||
def test_repr(self):
|
||||
items = pdf_inspector.extract_text_with_positions(
|
||||
fixture_path("thermo-freon12.pdf")
|
||||
)
|
||||
r = repr(items[0])
|
||||
assert "TextItem" in r
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Error handling
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestErrors:
|
||||
def test_nonexistent_file(self):
|
||||
with pytest.raises(ValueError):
|
||||
pdf_inspector.process_pdf("/nonexistent/file.pdf")
|
||||
|
||||
def test_not_a_pdf(self):
|
||||
with pytest.raises(ValueError):
|
||||
pdf_inspector.process_pdf_bytes(b"this is not a pdf")
|
||||
|
||||
def test_empty_bytes(self):
|
||||
with pytest.raises(ValueError):
|
||||
pdf_inspector.process_pdf_bytes(b"")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Multiple fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMultipleFixtures:
|
||||
"""Run basic processing on all available test fixtures."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename",
|
||||
[f for f in os.listdir(FIXTURES_DIR) if f.endswith(".pdf")],
|
||||
)
|
||||
def test_process_all_fixtures(self, filename):
|
||||
result = pdf_inspector.process_pdf(fixture_path(filename))
|
||||
assert result.pdf_type in (
|
||||
"text_based",
|
||||
"scanned",
|
||||
"image_based",
|
||||
"mixed",
|
||||
)
|
||||
assert result.page_count > 0
|
||||
assert result.confidence >= 0.0
|
||||
Reference in New Issue
Block a user