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, 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 cell_boxes: Option<Vec<Vec<Option<[f32; 4]>>>>,
}
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 {
match self.rows.get_mut(row).and_then(|r| r.get_mut(col)) {
Some(cell) => {
*cell = text.into();
true
}
None => false,
}
}
pub fn cell_bbox(&self, row: usize, col: usize) -> Option<[f32; 4]> {
*self.cell_boxes.as_ref()?.get(row)?.get(col)?
}
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 boxes = self
.cell_boxes
.get_or_insert_with(|| self.rows.iter().map(|r| vec![None; r.len()]).collect());
while boxes.len() < row + 1 {
boxes.push(Vec::new());
}
let brow = &mut boxes[row];
while brow.len() < col + 1 {
brow.push(None);
}
brow[col] = Some(bbox);
true
}
pub fn find_cell_by_bbox(&self, bbox: [f32; 4]) -> Option<(usize, usize)> {
let boxes = self.cell_boxes.as_ref()?;
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 (r, row) in boxes.iter().enumerate() {
for (c, cell) in row.iter().enumerate() {
let Some(cb) = cell 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, (r, c)));
}
}
}
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 table() -> Table {
Table {
rows: vec![
vec!["Year".into(), "Ducks".into()],
vec!["2019".into(), "120".into()],
],
cell_boxes: Some(vec![
vec![Some([0.0, 0.0, 50.0, 10.0]), Some([50.0, 0.0, 100.0, 10.0])],
vec![
Some([0.0, 10.0, 50.0, 20.0]),
Some([50.0, 10.0, 100.0, 20.0]),
],
]),
..Default::default()
}
}
#[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)));
}
}