use crate::model::{Paragraph, Table};
#[derive(Debug, Clone, Default)]
pub enum VisitorAction {
#[default]
Continue,
Replace(String),
Skip,
}
impl VisitorAction {
pub fn should_skip(&self) -> bool {
matches!(self, VisitorAction::Skip)
}
pub fn is_replace(&self) -> bool {
matches!(self, VisitorAction::Replace(_))
}
pub fn replacement(&self) -> Option<&str> {
match self {
VisitorAction::Replace(s) => Some(s),
_ => None,
}
}
}
pub trait DocumentVisitor: Send + Sync {
fn visit_paragraph(&mut self, para: &Paragraph) -> VisitorAction {
let _ = para;
VisitorAction::Continue
}
fn visit_table(&mut self, table: &Table) -> VisitorAction {
let _ = table;
VisitorAction::Continue
}
fn visit_image(&mut self, id: &str, alt: Option<&str>) -> VisitorAction {
let _ = (id, alt);
VisitorAction::Continue
}
fn visit_heading(&mut self, text: &str, level: u8) -> VisitorAction {
let _ = (text, level);
VisitorAction::Continue
}
fn visit_horizontal_rule(&mut self) -> VisitorAction {
VisitorAction::Continue
}
fn visit_list_item(&mut self, para: &Paragraph, level: u8, ordered: bool) -> VisitorAction {
let _ = (para, level, ordered);
VisitorAction::Continue
}
fn visit_raw(&mut self, content: &str) -> VisitorAction {
let _ = content;
VisitorAction::Continue
}
fn on_page_start(&mut self, page_number: u32) {
let _ = page_number;
}
fn on_page_end(&mut self, page_number: u32) {
let _ = page_number;
}
}
#[derive(Debug, Clone, Default)]
pub struct DefaultVisitor;
impl DefaultVisitor {
pub fn new() -> Self {
Self
}
}
impl DocumentVisitor for DefaultVisitor {}
#[derive(Debug, Clone, Default)]
pub struct SkipImagesVisitor;
impl DocumentVisitor for SkipImagesVisitor {
fn visit_image(&mut self, _id: &str, _alt: Option<&str>) -> VisitorAction {
VisitorAction::Skip
}
}
#[derive(Debug, Clone, Default)]
pub struct SimpleTableVisitor;
impl DocumentVisitor for SimpleTableVisitor {
fn visit_table(&mut self, table: &Table) -> VisitorAction {
let mut output = String::new();
for row in &table.rows {
let cells: Vec<String> = row.cells.iter().map(|c| c.plain_text()).collect();
output.push_str(&cells.join(" | "));
output.push('\n');
}
output.push('\n');
VisitorAction::Replace(output)
}
}
#[derive(Debug, Clone)]
pub struct MaxHeadingDepthVisitor {
max_level: u8,
}
impl MaxHeadingDepthVisitor {
pub fn new(max_level: u8) -> Self {
Self {
max_level: max_level.clamp(1, 6),
}
}
}
impl DocumentVisitor for MaxHeadingDepthVisitor {
fn visit_heading(&mut self, text: &str, level: u8) -> VisitorAction {
let effective_level = level.min(self.max_level);
let prefix = "#".repeat(effective_level as usize);
VisitorAction::Replace(format!("{} {}\n\n", prefix, text))
}
}
pub struct CompositeVisitor {
visitors: Vec<Box<dyn DocumentVisitor>>,
}
impl CompositeVisitor {
pub fn new() -> Self {
Self {
visitors: Vec::new(),
}
}
pub fn with_visitor<V: DocumentVisitor + 'static>(mut self, visitor: V) -> Self {
self.visitors.push(Box::new(visitor));
self
}
}
impl Default for CompositeVisitor {
fn default() -> Self {
Self::new()
}
}
impl DocumentVisitor for CompositeVisitor {
fn visit_paragraph(&mut self, para: &Paragraph) -> VisitorAction {
for visitor in &mut self.visitors {
let action = visitor.visit_paragraph(para);
if !matches!(action, VisitorAction::Continue) {
return action;
}
}
VisitorAction::Continue
}
fn visit_table(&mut self, table: &Table) -> VisitorAction {
for visitor in &mut self.visitors {
let action = visitor.visit_table(table);
if !matches!(action, VisitorAction::Continue) {
return action;
}
}
VisitorAction::Continue
}
fn visit_image(&mut self, id: &str, alt: Option<&str>) -> VisitorAction {
for visitor in &mut self.visitors {
let action = visitor.visit_image(id, alt);
if !matches!(action, VisitorAction::Continue) {
return action;
}
}
VisitorAction::Continue
}
fn visit_heading(&mut self, text: &str, level: u8) -> VisitorAction {
for visitor in &mut self.visitors {
let action = visitor.visit_heading(text, level);
if !matches!(action, VisitorAction::Continue) {
return action;
}
}
VisitorAction::Continue
}
fn visit_horizontal_rule(&mut self) -> VisitorAction {
for visitor in &mut self.visitors {
let action = visitor.visit_horizontal_rule();
if !matches!(action, VisitorAction::Continue) {
return action;
}
}
VisitorAction::Continue
}
fn visit_list_item(&mut self, para: &Paragraph, level: u8, ordered: bool) -> VisitorAction {
for visitor in &mut self.visitors {
let action = visitor.visit_list_item(para, level, ordered);
if !matches!(action, VisitorAction::Continue) {
return action;
}
}
VisitorAction::Continue
}
fn on_page_start(&mut self, page_number: u32) {
for visitor in &mut self.visitors {
visitor.on_page_start(page_number);
}
}
fn on_page_end(&mut self, page_number: u32) {
for visitor in &mut self.visitors {
visitor.on_page_end(page_number);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_visitor_action_default() {
let action = VisitorAction::default();
assert!(matches!(action, VisitorAction::Continue));
}
#[test]
fn test_visitor_action_should_skip() {
assert!(!VisitorAction::Continue.should_skip());
assert!(!VisitorAction::Replace("test".into()).should_skip());
assert!(VisitorAction::Skip.should_skip());
}
#[test]
fn test_visitor_action_replacement() {
assert!(VisitorAction::Continue.replacement().is_none());
assert!(VisitorAction::Skip.replacement().is_none());
assert_eq!(
VisitorAction::Replace("hello".into()).replacement(),
Some("hello")
);
}
#[test]
fn test_default_visitor() {
let mut visitor = DefaultVisitor::new();
let para = Paragraph::new();
let action = visitor.visit_paragraph(¶);
assert!(matches!(action, VisitorAction::Continue));
}
#[test]
fn test_skip_images_visitor() {
let mut visitor = SkipImagesVisitor;
let action = visitor.visit_image("img1", Some("alt text"));
assert!(action.should_skip());
}
#[test]
fn test_max_heading_depth_visitor() {
let mut visitor = MaxHeadingDepthVisitor::new(2);
let action = visitor.visit_heading("Deep Heading", 4);
assert!(action.is_replace());
let replacement = action.replacement().unwrap();
assert!(replacement.starts_with("## "));
}
#[test]
fn test_composite_visitor() {
let mut composite = CompositeVisitor::new()
.with_visitor(SkipImagesVisitor)
.with_visitor(DefaultVisitor);
let action = composite.visit_image("img1", None);
assert!(action.should_skip());
let para = Paragraph::new();
let action = composite.visit_paragraph(¶);
assert!(matches!(action, VisitorAction::Continue));
}
}