1use std::cmp::Ordering;
2
3#[cfg(feature = "semantic")]
4use serde::{Deserialize, Serialize};
5
6#[derive(Debug, Clone)]
11#[cfg_attr(
12 feature = "semantic",
13 derive(Serialize, Deserialize),
14 serde(tag = "type", rename_all = "snake_case")
15)]
16pub enum Element {
17 Title(ElementData),
18 Paragraph(ElementData),
19 Table(TableElementData),
20 Header(ElementData),
21 Footer(ElementData),
22 ListItem(ElementData),
23 Image(ImageElementData),
24 CodeBlock(ElementData),
25 KeyValue(KeyValueElementData),
26}
27
28impl Element {
29 pub fn text(&self) -> &str {
31 match self {
32 Self::Title(d)
33 | Self::Paragraph(d)
34 | Self::Header(d)
35 | Self::Footer(d)
36 | Self::ListItem(d)
37 | Self::CodeBlock(d) => &d.text,
38 Self::Table(t) => {
39 let _ = t;
42 ""
43 }
44 Self::Image(img) => img.alt_text.as_deref().unwrap_or(""),
45 Self::KeyValue(kv) => &kv.value,
46 }
47 }
48
49 pub fn display_text(&self) -> String {
57 match self {
58 Self::Table(t) => t
59 .rows
60 .iter()
61 .map(|row| row.join(" | "))
62 .collect::<Vec<_>>()
63 .join("\n"),
64 Self::Image(img) => img.alt_text.clone().unwrap_or_default(),
65 Self::KeyValue(kv) => format!("{}: {}", kv.key, kv.value),
66 _ => self.text().to_string(),
67 }
68 }
69
70 pub fn page(&self) -> u32 {
72 self.metadata().page
73 }
74
75 pub fn bbox(&self) -> &ElementBBox {
77 &self.metadata().bbox
78 }
79
80 pub fn metadata(&self) -> &ElementMetadata {
82 match self {
83 Self::Title(d)
84 | Self::Paragraph(d)
85 | Self::Header(d)
86 | Self::Footer(d)
87 | Self::ListItem(d)
88 | Self::CodeBlock(d) => &d.metadata,
89 Self::Table(t) => &t.metadata,
90 Self::Image(img) => &img.metadata,
91 Self::KeyValue(kv) => &kv.metadata,
92 }
93 }
94
95 pub fn metadata_mut(&mut self) -> &mut ElementMetadata {
97 match self {
98 Self::Title(d)
99 | Self::Paragraph(d)
100 | Self::Header(d)
101 | Self::Footer(d)
102 | Self::ListItem(d)
103 | Self::CodeBlock(d) => &mut d.metadata,
104 Self::Table(t) => &mut t.metadata,
105 Self::Image(img) => &mut img.metadata,
106 Self::KeyValue(kv) => &mut kv.metadata,
107 }
108 }
109
110 pub fn type_name(&self) -> &'static str {
116 match self {
117 Self::Title(_) => "title",
118 Self::Paragraph(_) => "paragraph",
119 Self::Table(_) => "table",
120 Self::Header(_) => "header",
121 Self::Footer(_) => "footer",
122 Self::ListItem(_) => "list_item",
123 Self::Image(_) => "image",
124 Self::CodeBlock(_) => "code_block",
125 Self::KeyValue(_) => "key_value",
126 }
127 }
128
129 pub fn set_parent_heading(&mut self, heading: Option<String>) {
131 self.metadata_mut().parent_heading = heading;
132 }
133
134 pub fn set_heading_path(&mut self, path: Vec<String>) {
136 self.metadata_mut().heading_path = path;
137 }
138
139 pub fn row_count(&self) -> Option<usize> {
141 match self {
142 Self::Table(t) => Some(t.rows.len()),
143 _ => None,
144 }
145 }
146
147 pub fn column_count(&self) -> Option<usize> {
149 match self {
150 Self::Table(t) => t.rows.first().map(|r| r.len()),
151 _ => None,
152 }
153 }
154
155 pub fn cell(&self, row: usize, col: usize) -> Option<&str> {
157 match self {
158 Self::Table(t) => t.rows.get(row).and_then(|r| r.get(col)).map(|s| s.as_str()),
159 _ => None,
160 }
161 }
162}
163
164pub fn element_reading_order(a: &Element, b: &Element) -> Ordering {
169 let page_cmp = a.page().cmp(&b.page());
170 if page_cmp != Ordering::Equal {
171 return page_cmp;
172 }
173 let y_cmp = b.bbox().y.total_cmp(&a.bbox().y);
175 if y_cmp != Ordering::Equal {
176 return y_cmp;
177 }
178 a.bbox().x.total_cmp(&b.bbox().x)
179}
180
181impl PartialEq for Element {
182 fn eq(&self, other: &Self) -> bool {
183 std::mem::discriminant(self) == std::mem::discriminant(other) && self.content_eq(other)
184 }
185}
186
187impl Eq for Element {}
188
189impl Element {
190 fn content_eq(&self, other: &Self) -> bool {
192 match (self, other) {
193 (Self::Title(a), Self::Title(b))
194 | (Self::Paragraph(a), Self::Paragraph(b))
195 | (Self::Header(a), Self::Header(b))
196 | (Self::Footer(a), Self::Footer(b))
197 | (Self::ListItem(a), Self::ListItem(b))
198 | (Self::CodeBlock(a), Self::CodeBlock(b)) => a.text == b.text,
199 (Self::Table(a), Self::Table(b)) => a.rows == b.rows,
200 (Self::Image(a), Self::Image(b)) => a.alt_text == b.alt_text,
201 (Self::KeyValue(a), Self::KeyValue(b)) => a.key == b.key && a.value == b.value,
202 _ => false,
203 }
204 }
205}
206
207#[derive(Debug, Clone)]
209#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
210pub struct ElementData {
211 pub text: String,
212 pub metadata: ElementMetadata,
213}
214
215#[derive(Debug, Clone)]
218#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
219pub struct RichCell {
220 pub row: usize,
221 pub col: usize,
222 pub row_span: usize,
223 pub col_span: usize,
224 pub text: String,
225 pub is_header: bool,
226}
227
228#[derive(Debug, Clone)]
232#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
233pub struct TableStructure {
234 pub cells: Vec<RichCell>,
235 pub num_rows: usize,
236 pub num_cols: usize,
237 pub header_rows: usize,
240}
241
242#[derive(Debug, Clone)]
244#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
245#[non_exhaustive]
246pub struct TableElementData {
247 pub rows: Vec<Vec<String>>,
251 pub structure: Option<TableStructure>,
253 pub metadata: ElementMetadata,
254}
255
256impl TableElementData {
257 pub fn new(rows: Vec<Vec<String>>, metadata: ElementMetadata) -> Self {
260 Self {
261 rows,
262 structure: None,
263 metadata,
264 }
265 }
266
267 pub fn from_structure(structure: TableStructure, metadata: ElementMetadata) -> Self {
270 let mut rows = vec![vec![String::new(); structure.num_cols]; structure.num_rows];
271 for cell in &structure.cells {
272 for r in cell.row..(cell.row + cell.row_span).min(structure.num_rows) {
273 for c in cell.col..(cell.col + cell.col_span).min(structure.num_cols) {
274 rows[r][c] = cell.text.clone();
275 }
276 }
277 }
278 Self {
279 rows,
280 structure: Some(structure),
281 metadata,
282 }
283 }
284}
285
286#[derive(Debug, Clone)]
288#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
289pub struct ImageElementData {
290 pub alt_text: Option<String>,
291 pub metadata: ElementMetadata,
292}
293
294#[derive(Debug, Clone)]
296#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
297pub struct KeyValueElementData {
298 pub key: String,
299 pub value: String,
300 pub metadata: ElementMetadata,
301}
302
303#[derive(Debug, Clone)]
305#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
306pub struct ElementMetadata {
307 pub page: u32,
309 pub bbox: ElementBBox,
311 pub confidence: f64,
313 pub font_name: Option<String>,
315 pub font_size: Option<f64>,
317 pub is_bold: bool,
319 pub is_italic: bool,
321 pub parent_heading: Option<String>,
323 pub heading_path: Vec<String>,
325 #[cfg(feature = "unstable-spi")]
330 #[cfg_attr(
331 feature = "semantic",
332 serde(default, skip_serializing_if = "Option::is_none")
333 )]
334 pub class_label: Option<String>,
335}
336
337impl Default for ElementMetadata {
338 fn default() -> Self {
339 Self {
340 page: 0,
341 bbox: ElementBBox::ZERO,
342 confidence: 1.0,
343 font_name: None,
344 font_size: None,
345 is_bold: false,
346 is_italic: false,
347 parent_heading: None,
348 heading_path: Vec::new(),
349 #[cfg(feature = "unstable-spi")]
350 class_label: None,
351 }
352 }
353}
354
355#[derive(Debug, Clone, Copy, PartialEq)]
357#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
358pub struct ElementBBox {
359 pub x: f64,
361 pub y: f64,
363 pub width: f64,
365 pub height: f64,
367}
368
369impl ElementBBox {
370 pub const ZERO: Self = Self {
372 x: 0.0,
373 y: 0.0,
374 width: 0.0,
375 height: 0.0,
376 };
377
378 pub fn new(x: f64, y: f64, width: f64, height: f64) -> Self {
380 Self {
381 x,
382 y,
383 width,
384 height,
385 }
386 }
387
388 pub fn right(&self) -> f64 {
390 self.x + self.width
391 }
392
393 pub fn top(&self) -> f64 {
395 self.y + self.height
396 }
397}