snapshot/
snapshot.rs

1//! Generate the PDF snapshot baseline.
2//!
3//! `cargo run --release -p docling-pdf --example snapshot -- <root> <outdir>`
4//!
5//! Recursively converts every `*.pdf` under `<root>` and writes its Markdown to
6//! `<outdir>/<relative-path>.md` (the pipeline's deterministic output, committed
7//! as the conformance baseline). Conversion errors are written verbatim as
8//! `ERROR: …` so failures are captured, never silently skipped.
9
10use std::path::{Path, PathBuf};
11
12use docling_pdf::Pipeline;
13
14const IMAGE_EXTS: &[&str] = &["png", "jpg", "jpeg", "tif", "tiff", "bmp", "gif", "webp"];
15
16fn is_supported(p: &Path) -> bool {
17    matches!(
18        p.extension().and_then(|e| e.to_str()),
19        Some(e) if e == "pdf" || e == "gz" || IMAGE_EXTS.contains(&e)
20    )
21}
22
23fn find_pdfs(dir: &Path, out: &mut Vec<PathBuf>) {
24    let Ok(entries) = std::fs::read_dir(dir) else {
25        return;
26    };
27    let mut entries: Vec<_> = entries.flatten().map(|e| e.path()).collect();
28    entries.sort();
29    for p in entries {
30        if p.is_dir() {
31            // Skip `large/` — big perf-test inputs with no conformance baseline —
32            // and `*_artifacts/` — gitignored by-products of `--images referenced`
33            // / dclx conversions that would otherwise be swept in as image inputs
34            // and mint snapshots that don't exist on a clean checkout.
35            if p.file_name()
36                .is_some_and(|n| n == "large" || n.to_string_lossy().ends_with("_artifacts"))
37            {
38                continue;
39            }
40            find_pdfs(&p, out);
41        } else if is_supported(&p) {
42            out.push(p);
43        }
44    }
45}
46
47fn main() {
48    let mut args = std::env::args().skip(1);
49    let root = PathBuf::from(args.next().expect("usage: snapshot <root> <outdir>"));
50    let outdir = PathBuf::from(args.next().expect("usage: snapshot <root> <outdir>"));
51
52    let mut pdfs = Vec::new();
53    find_pdfs(&root, &mut pdfs);
54
55    let mut pipeline = Pipeline::new().expect("load pipeline");
56    let (mut ok, mut err) = (0u32, 0u32);
57    for pdf in &pdfs {
58        let rel = pdf.strip_prefix(&root).unwrap_or(pdf);
59        let name = pdf.file_name().unwrap().to_string_lossy().to_string();
60        let md = match std::fs::read(pdf)
61            .map_err(|e| format!("read: {e}"))
62            .and_then(|bytes| {
63                let ext = pdf.extension().and_then(|e| e.to_str()).unwrap_or("");
64                let result = if ext == "gz" {
65                    docling_pdf::convert_mets_gbs(&bytes, &name)
66                } else if IMAGE_EXTS.contains(&ext) {
67                    pipeline.convert_image(&bytes, &name)
68                } else {
69                    pipeline.convert(&bytes, None, &name)
70                };
71                result
72                    .map(|d| d.export_to_markdown())
73                    .map_err(|e| e.to_string())
74            }) {
75            Ok(md) => {
76                ok += 1;
77                md
78            }
79            Err(e) => {
80                err += 1;
81                eprintln!("ERR {}: {e}", rel.display());
82                format!("ERROR: {e}\n")
83            }
84        };
85        // Groundtruth naming keeps the source extension: `<file>.<ext>.md`.
86        let mut dest = outdir.join(rel).into_os_string();
87        dest.push(".md");
88        let dest = PathBuf::from(dest);
89        std::fs::create_dir_all(dest.parent().unwrap()).expect("mkdir");
90        std::fs::write(&dest, md).expect("write snapshot");
91    }
92    eprintln!("snapshots: {} ok, {} error, {} total", ok, err, pdfs.len());
93}
snapshot/snapshot.rs

snapshot/
snapshot.rs
