use crate::markdown::{to_markdown, to_markdown_images};
use crate::ImageMode;
#[derive(Debug, Clone, PartialEq)]
pub struct DoclingDocument {
pub name: String,
pub nodes: Vec<Node>,
pub strict_markdown: bool,
pub compact_tables: bool,
pub links: Vec<(String, String)>,
pub confidence: Option<crate::confidence::ConfidenceReport>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Node {
Heading { level: u8, text: String },
Paragraph { text: String },
CheckboxItem { checked: bool, text: String },
ListItem {
ordered: bool,
number: u64,
first_in_list: bool,
text: String,
level: u8,
marker: Option<String>,
location: Option<[u16; 4]>,
dclx: Option<ListItemDclx>,
href: Option<String>,
layer: Option<ContentLayer>,
},
Code {
language: Option<String>,
text: String,
orig: Option<String>,
pretty: Option<String>,
},
Table(Table),
Picture {
caption: Option<String>,
image: Option<PictureImage>,
classification: Option<Vec<PictureClass>>,
},
Formula {
latex: String,
orig: String,
location: Option<[u16; 4]>,
},
Chart {
kind: String,
table: Table,
caption: Option<String>,
location: Option<[u16; 4]>,
},
Group { label: String, children: Vec<Node> },
FieldRegion { items: Vec<FieldItem> },
InlineGroup {
unwrapped: bool,
runs: Vec<InlineRun>,
md_text: String,
},
Furniture {
layer: ContentLayer,
inner: Box<Node>,
},
Located {
location: [u16; 4],
inner: Box<Node>,
},
PageFurniture {
footer: bool,
location: [u16; 4],
text: String,
},
PageBreak,
PageInfo {
page_no: usize,
width: f32,
height: f32,
},
DoclangOnly(Box<Node>),
TextDump(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Script {
#[default]
Baseline,
Sub,
Super,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct InlineRun {
pub text: String,
pub bold: bool,
pub italic: bool,
pub underline: bool,
pub strike: bool,
pub script: Script,
pub code: bool,
pub formula: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ContentLayer {
Furniture,
Notes,
Invisible,
}
impl ContentLayer {
pub fn value(self) -> &'static str {
match self {
ContentLayer::Furniture => "furniture",
ContentLayer::Notes => "notes",
ContentLayer::Invisible => "invisible",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct ListItemDclx {
pub ordered: bool,
pub marker: Option<String>,
pub text: String,
pub runs: Vec<InlineRun>,
}
impl InlineRun {
pub fn is_plain(&self) -> bool {
!self.bold
&& !self.italic
&& !self.underline
&& !self.strike
&& !self.code
&& !self.formula
&& self.script == Script::Baseline
}
}
pub fn inline_paragraph_node(md_text: String, runs: Vec<InlineRun>, unwrapped: bool) -> Node {
let single_plain = runs.len() <= 1 && runs.first().is_none_or(|r| r.is_plain());
if single_plain {
Node::Paragraph { text: md_text }
} else {
Node::InlineGroup {
unwrapped: unwrapped && runs.len() >= 2,
runs,
md_text,
}
}
}
#[derive(Debug, Clone, PartialEq, Default)]
pub struct FieldItem {
pub marker: Option<String>,
pub key: Option<String>,
pub value: Option<String>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct PictureClass {
pub class_name: String,
pub confidence: f32,
}
#[derive(Debug, Clone, PartialEq)]
pub struct PictureImage {
pub mimetype: String,
pub width: u32,
pub height: u32,
pub data: Vec<u8>,
}
impl PictureImage {
pub fn data_uri(&self) -> String {
format!(
"data:{};base64,{}",
self.mimetype,
crate::base64::encode(&self.data)
)
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct TableCell {
pub text: String,
pub bbox: Option<[f32; 4]>,
pub start_row: usize,
pub start_col: usize,
pub row_span: usize,
pub col_span: usize,
pub column_header: bool,
pub row_header: bool,
pub row_section: bool,
}
#[derive(Debug, Clone, PartialEq, Default)]
pub struct Table {
pub rows: Vec<Vec<String>>,
pub location: Option<[u16; 4]>,
pub structure: Option<TableStructure>,
pub cell_blocks: Option<Vec<Vec<Vec<Node>>>>,
pub caption: Option<String>,
pub cells: Option<Vec<TableCell>>,
}
impl Table {
pub fn cell_text(&self, row: usize, col: usize) -> Option<&str> {
self.rows.get(row)?.get(col).map(String::as_str)
}
pub fn set_cell_text(&mut self, row: usize, col: usize, text: impl Into<String>) -> bool {
if self.rows.get(row).and_then(|r| r.get(col)).is_none() {
return false;
}
let text = text.into();
let covering = self.cells.as_mut().and_then(|cells| {
cells.iter_mut().find(|c| {
(c.start_row..c.start_row + c.row_span).contains(&row)
&& (c.start_col..c.start_col + c.col_span).contains(&col)
})
});
if let Some(cell) = covering {
cell.text = text.clone();
let (r0, r1) = (cell.start_row, cell.start_row + cell.row_span);
let (c0, c1) = (cell.start_col, cell.start_col + cell.col_span);
for r in self.rows.iter_mut().take(r1).skip(r0) {
for slot in r.iter_mut().take(c1).skip(c0) {
*slot = text.clone();
}
}
} else {
self.rows[row][col] = text;
}
true
}
pub fn derive_cells(&self) -> Vec<TableCell> {
let s = self.structure.as_ref();
let flag = |grid: Option<&Vec<Vec<bool>>>, r: usize, c: usize| {
grid.and_then(|g| g.get(r))
.and_then(|row| row.get(c))
.copied()
.unwrap_or(false)
};
let col_cont = |r: usize, c: usize| flag(s.map(|s| &s.col_continuation), r, c);
let row_cont = |r: usize, c: usize| flag(s.map(|s| &s.row_continuation), r, c);
let is_col_header = |r: usize, c: usize| match s {
Some(st) if !st.col_header.is_empty() => flag(Some(&st.col_header), r, c),
Some(st) if !st.header_row.is_empty() => st.header_row.get(r).copied().unwrap_or(false),
_ => r == 0,
};
let mut cells = Vec::new();
for (r, row) in self.rows.iter().enumerate() {
for (c, text) in row.iter().enumerate() {
if col_cont(r, c) || row_cont(r, c) {
continue; }
let mut col_span = 1;
while c + col_span < row.len() && col_cont(r, c + col_span) {
col_span += 1;
}
let mut row_span = 1;
while r + row_span < self.rows.len() && row_cont(r + row_span, c) {
row_span += 1;
}
cells.push(TableCell {
text: text.clone(),
bbox: None,
start_row: r,
start_col: c,
row_span,
col_span,
column_header: is_col_header(r, c),
row_header: flag(s.map(|s| &s.row_header), r, c),
row_section: false,
});
}
}
cells
}
pub fn cell_at(&self, row: usize, col: usize) -> Option<&TableCell> {
self.cells.as_ref()?.iter().find(|c| {
(c.start_row..c.start_row + c.row_span).contains(&row)
&& (c.start_col..c.start_col + c.col_span).contains(&col)
})
}
pub fn cell_bbox(&self, row: usize, col: usize) -> Option<[f32; 4]> {
self.cell_at(row, col)?.bbox
}
pub fn set_cell_bbox(&mut self, row: usize, col: usize, bbox: [f32; 4]) -> bool {
if self.rows.get(row).and_then(|r| r.get(col)).is_none() {
return false;
}
let rows = &self.rows;
let cells = self.cells.get_or_insert_with(|| {
rows.iter()
.enumerate()
.flat_map(|(r, cols)| {
cols.iter().enumerate().map(move |(c, text)| TableCell {
text: text.clone(),
bbox: None,
start_row: r,
start_col: c,
row_span: 1,
col_span: 1,
column_header: false,
row_header: false,
row_section: false,
})
})
.collect()
});
match cells.iter_mut().find(|c| {
(c.start_row..c.start_row + c.row_span).contains(&row)
&& (c.start_col..c.start_col + c.col_span).contains(&col)
}) {
Some(cell) => {
cell.bbox = Some(bbox);
true
}
None => {
cells.push(TableCell {
text: self.rows[row][col].clone(),
bbox: Some(bbox),
start_row: row,
start_col: col,
row_span: 1,
col_span: 1,
column_header: false,
row_header: false,
row_section: false,
});
true
}
}
}
pub fn find_cell_by_bbox(&self, bbox: [f32; 4]) -> Option<(usize, usize)> {
let area = |b: &[f32; 4]| ((b[2] - b[0]) * (b[3] - b[1])).max(0.0);
let mut best: Option<(f32, (usize, usize))> = None;
for cell in self.cells.as_deref()?.iter() {
let Some(cb) = cell.bbox else { continue };
let iw = (bbox[2].min(cb[2]) - bbox[0].max(cb[0])).max(0.0);
let ih = (bbox[3].min(cb[3]) - bbox[1].max(cb[1])).max(0.0);
let inter = iw * ih;
if inter <= 0.0 {
continue;
}
let iou = inter / (area(&bbox) + area(&cb) - inter).max(f32::EPSILON);
if best.is_none_or(|(b, _)| iou > b) {
best = Some((iou, (cell.start_row, cell.start_col)));
}
}
best.map(|(_, pos)| pos)
}
pub fn update_cell_by_bbox(
&mut self,
bbox: [f32; 4],
text: impl Into<String>,
) -> Option<(usize, usize)> {
let (row, col) = self.find_cell_by_bbox(bbox)?;
self.set_cell_text(row, col, text);
Some((row, col))
}
}
#[derive(Debug, Clone, PartialEq, Default)]
pub struct TableStructure {
pub header_row: Vec<bool>,
pub col_continuation: Vec<Vec<bool>>,
pub row_continuation: Vec<Vec<bool>>,
pub row_header: Vec<Vec<bool>>,
pub col_header: Vec<Vec<bool>>,
}
impl DoclingDocument {
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
nodes: Vec::new(),
strict_markdown: false,
compact_tables: false,
links: Vec::new(),
confidence: None,
}
}
pub fn tables(&self) -> impl Iterator<Item = &Table> {
fn unwrap_table(n: &Node) -> Option<&Table> {
match n {
Node::Table(t) => Some(t),
Node::Located { inner, .. } => unwrap_table(inner),
_ => None,
}
}
self.nodes.iter().filter_map(unwrap_table)
}
pub fn tables_mut(&mut self) -> impl Iterator<Item = &mut Table> {
fn unwrap_table(n: &mut Node) -> Option<&mut Table> {
match n {
Node::Table(t) => Some(t),
Node::Located { inner, .. } => unwrap_table(inner),
_ => None,
}
}
self.nodes.iter_mut().filter_map(unwrap_table)
}
pub fn push(&mut self, node: Node) {
self.nodes.push(node);
}
pub fn add_heading(&mut self, level: u8, text: impl Into<String>) {
self.push(Node::Heading {
level,
text: text.into(),
});
}
pub fn add_paragraph(&mut self, text: impl Into<String>) {
self.push(Node::Paragraph { text: text.into() });
}
pub fn export_to_markdown(&self) -> String {
to_markdown(self, self.strict_markdown)
}
pub fn export_to_markdown_with(&self, strict: bool) -> String {
to_markdown(self, strict)
}
pub fn export_to_json(&self) -> String {
serde_json::to_string_pretty(&self.export_to_json_value())
.expect("DoclingDocument JSON is always serializable")
}
pub fn export_to_json_value(&self) -> serde_json::Value {
crate::json::to_json(self)
}
pub fn export_to_doclang(&self) -> String {
crate::doclang::export_to_doclang(&self.nodes)
}
pub fn export_to_markdown_with_images(
&self,
image_mode: ImageMode,
artifacts_dir: &str,
) -> (String, Vec<(String, Vec<u8>)>) {
to_markdown_images(self, self.strict_markdown, image_mode, artifacts_dir)
}
}
#[cfg(test)]
mod table_api_tests {
use super::*;
fn cell(
text: &str,
bbox: [f32; 4],
(start_row, start_col): (usize, usize),
(row_span, col_span): (usize, usize),
) -> TableCell {
TableCell {
text: text.into(),
bbox: Some(bbox),
start_row,
start_col,
row_span,
col_span,
column_header: false,
row_header: false,
row_section: false,
}
}
fn table() -> Table {
Table {
rows: vec![
vec!["Year".into(), "Ducks".into()],
vec!["2019".into(), "120".into()],
],
cells: Some(vec![
cell("Year", [0.0, 0.0, 50.0, 10.0], (0, 0), (1, 1)),
cell("Ducks", [50.0, 0.0, 100.0, 10.0], (0, 1), (1, 1)),
cell("2019", [0.0, 10.0, 50.0, 20.0], (1, 0), (1, 1)),
cell("120", [50.0, 10.0, 100.0, 20.0], (1, 1), (1, 1)),
]),
..Default::default()
}
}
#[test]
fn derive_cells_reads_spans_and_headers_from_structure() {
let t = Table {
rows: vec![
vec!["Wide".into(), "Wide".into(), "C".into()],
vec!["a".into(), "b".into(), "c".into()],
],
structure: Some(TableStructure {
header_row: vec![true, false],
col_continuation: vec![vec![false, true, false], vec![false; 3]],
row_continuation: vec![vec![false; 3], vec![false; 3]],
row_header: Vec::new(),
col_header: Vec::new(),
}),
..Default::default()
};
let cells = t.derive_cells();
assert_eq!(cells.len(), 5, "two anchors in row 0, three in row 1");
let wide = &cells[0];
assert_eq!((wide.col_span, wide.row_span), (2, 1));
assert!(wide.column_header, "header_row band");
assert!(cells.iter().skip(2).all(|c| !c.column_header));
let plain = Table {
rows: vec![vec!["h".into()], vec!["x".into()]],
..Default::default()
};
let cells = plain.derive_cells();
assert_eq!(cells.len(), 2);
assert!(cells[0].column_header && !cells[1].column_header);
}
#[test]
fn span_repair_updates_the_whole_cell() {
let mut t = Table {
rows: vec![
vec!["Wide".into(), "Wide".into(), "C".into()],
vec!["a".into(), "b".into(), "c".into()],
],
cells: Some(vec![
cell("Wide", [0.0, 0.0, 100.0, 10.0], (0, 0), (1, 2)),
cell("C", [100.0, 0.0, 150.0, 10.0], (0, 2), (1, 1)),
]),
..Default::default()
};
assert!(t.set_cell_text(0, 1, "Fixed"));
assert_eq!(
t.rows[0],
vec!["Fixed".to_string(), "Fixed".into(), "C".into()]
);
assert_eq!(t.cell_at(0, 1).unwrap().text, "Fixed");
assert_eq!(t.cell_at(0, 1).unwrap().col_span, 2);
}
#[test]
fn bbox_lookup_and_repair_flow_into_exports() {
let mut doc = DoclingDocument::new("t");
doc.push(Node::Table(table()));
assert_eq!(doc.tables().count(), 1);
let t = doc.tables_mut().next().unwrap();
assert_eq!(t.find_cell_by_bbox([52.0, 11.0, 98.0, 19.0]), Some((1, 1)));
assert_eq!(
t.update_cell_by_bbox([52.0, 11.0, 98.0, 19.0], "125"),
Some((1, 1))
);
assert_eq!(t.cell_text(1, 1), Some("125"));
assert!(doc.export_to_markdown().contains("125"));
let t = doc.tables_mut().next().unwrap();
assert_eq!(t.find_cell_by_bbox([500.0, 500.0, 600.0, 600.0]), None);
}
#[test]
fn cell_accessors_bound_check_and_geometry_materializes() {
let mut t = table();
assert_eq!(t.cell_text(0, 0), Some("Year"));
assert_eq!(t.cell_text(5, 0), None);
assert!(!t.set_cell_text(0, 9, "x"), "outside the grid");
assert_eq!(t.cell_bbox(1, 0), Some([0.0, 10.0, 50.0, 20.0]));
let mut plain = Table {
rows: vec![vec!["a".into(), "b".into()]],
..Default::default()
};
assert_eq!(plain.cell_bbox(0, 1), None);
assert!(!plain.set_cell_bbox(0, 5, [0.0; 4]), "outside the grid");
assert!(plain.set_cell_bbox(0, 1, [1.0, 2.0, 3.0, 4.0]));
assert_eq!(plain.cell_bbox(0, 1), Some([1.0, 2.0, 3.0, 4.0]));
assert_eq!(plain.find_cell_by_bbox([1.5, 2.5, 2.5, 3.5]), Some((0, 1)));
}
}