markdown_writer.rs

  1use std::cell::RefCell;
  2use std::collections::VecDeque;
  3use std::rc::Rc;
  4use std::sync::OnceLock;
  5
  6use anyhow::Result;
  7use markup5ever_rcdom::{Handle, NodeData};
  8use regex::Regex;
  9
 10use crate::html_element::HtmlElement;
 11
 12fn empty_line_regex() -> &'static Regex {
 13    static REGEX: OnceLock<Regex> = OnceLock::new();
 14    REGEX.get_or_init(|| Regex::new(r"^\s*$").unwrap())
 15}
 16
 17fn more_than_three_newlines_regex() -> &'static Regex {
 18    static REGEX: OnceLock<Regex> = OnceLock::new();
 19    REGEX.get_or_init(|| Regex::new(r"\n{3,}").unwrap())
 20}
 21
 22pub enum StartTagOutcome {
 23    Continue,
 24    Skip,
 25}
 26
 27pub type TagHandler = Rc<RefCell<dyn HandleTag>>;
 28
 29pub struct MarkdownWriter {
 30    current_element_stack: VecDeque<HtmlElement>,
 31    pub(crate) markdown: String,
 32}
 33
 34impl MarkdownWriter {
 35    pub fn new() -> Self {
 36        Self {
 37            current_element_stack: VecDeque::new(),
 38            markdown: String::new(),
 39        }
 40    }
 41
 42    pub fn current_element_stack(&self) -> &VecDeque<HtmlElement> {
 43        &self.current_element_stack
 44    }
 45
 46    pub fn is_inside(&self, tag: &str) -> bool {
 47        self.current_element_stack
 48            .iter()
 49            .any(|parent_element| parent_element.tag == tag)
 50    }
 51
 52    /// Appends the given string slice onto the end of the Markdown output.
 53    pub fn push_str(&mut self, str: &str) {
 54        self.markdown.push_str(str);
 55    }
 56
 57    /// Appends a newline to the end of the Markdown output.
 58    pub fn push_newline(&mut self) {
 59        self.push_str("\n");
 60    }
 61
 62    /// Appends a blank line to the end of the Markdown output.
 63    pub fn push_blank_line(&mut self) {
 64        self.push_str("\n\n");
 65    }
 66
 67    pub fn run(mut self, root_node: &Handle, handlers: &mut Vec<TagHandler>) -> Result<String> {
 68        self.visit_node(&root_node, handlers)?;
 69        Ok(Self::prettify_markdown(self.markdown))
 70    }
 71
 72    fn prettify_markdown(markdown: String) -> String {
 73        let markdown = empty_line_regex().replace_all(&markdown, "");
 74        let markdown = more_than_three_newlines_regex().replace_all(&markdown, "\n\n");
 75
 76        markdown.trim().to_string()
 77    }
 78
 79    fn visit_node(&mut self, node: &Handle, handlers: &mut [TagHandler]) -> Result<()> {
 80        let mut current_element = None;
 81
 82        match node.data {
 83            NodeData::Document
 84            | NodeData::Doctype { .. }
 85            | NodeData::ProcessingInstruction { .. }
 86            | NodeData::Comment { .. } => {
 87                // Currently left unimplemented, as we're not interested in this data
 88                // at this time.
 89            }
 90            NodeData::Element {
 91                ref name,
 92                ref attrs,
 93                ..
 94            } => {
 95                let tag_name = name.local.to_string();
 96                if !tag_name.is_empty() {
 97                    current_element = Some(HtmlElement {
 98                        tag: tag_name,
 99                        attrs: attrs.clone(),
100                    });
101                }
102            }
103            NodeData::Text { ref contents } => {
104                let text = contents.borrow().to_string();
105                self.visit_text(text, handlers)?;
106            }
107        }
108
109        if let Some(current_element) = current_element.as_ref() {
110            match self.start_tag(&current_element, handlers) {
111                StartTagOutcome::Continue => {}
112                StartTagOutcome::Skip => return Ok(()),
113            }
114
115            self.current_element_stack
116                .push_back(current_element.clone());
117        }
118
119        for child in node.children.borrow().iter() {
120            self.visit_node(child, handlers)?;
121        }
122
123        if let Some(current_element) = current_element {
124            self.current_element_stack.pop_back();
125            self.end_tag(&current_element, handlers);
126        }
127
128        Ok(())
129    }
130
131    fn start_tag(&mut self, tag: &HtmlElement, handlers: &mut [TagHandler]) -> StartTagOutcome {
132        for handler in handlers {
133            if handler.borrow().should_handle(tag.tag.as_str()) {
134                match handler.borrow_mut().handle_tag_start(tag, self) {
135                    StartTagOutcome::Continue => {}
136                    StartTagOutcome::Skip => return StartTagOutcome::Skip,
137                }
138            }
139        }
140
141        StartTagOutcome::Continue
142    }
143
144    fn end_tag(&mut self, tag: &HtmlElement, handlers: &mut [TagHandler]) {
145        for handler in handlers {
146            if handler.borrow().should_handle(tag.tag.as_str()) {
147                handler.borrow_mut().handle_tag_end(tag, self);
148            }
149        }
150    }
151
152    fn visit_text(&mut self, text: String, handlers: &mut [TagHandler]) -> Result<()> {
153        for handler in handlers {
154            match handler.borrow_mut().handle_text(&text, self) {
155                HandlerOutcome::Handled => return Ok(()),
156                HandlerOutcome::NoOp => {}
157            }
158        }
159
160        let text = text
161            .trim_matches(|char| char == '\n' || char == '\r' || char == '\t')
162            .replace('\n', " ");
163
164        self.push_str(&text);
165
166        Ok(())
167    }
168}
169
170pub enum HandlerOutcome {
171    Handled,
172    NoOp,
173}
174
175pub trait HandleTag {
176    /// Returns whether this handler should handle the given tag.
177    fn should_handle(&self, tag: &str) -> bool;
178
179    /// Handles the start of the given tag.
180    fn handle_tag_start(
181        &mut self,
182        _tag: &HtmlElement,
183        _writer: &mut MarkdownWriter,
184    ) -> StartTagOutcome {
185        StartTagOutcome::Continue
186    }
187
188    /// Handles the end of the given tag.
189    fn handle_tag_end(&mut self, _tag: &HtmlElement, _writer: &mut MarkdownWriter) {}
190
191    fn handle_text(&mut self, _text: &str, _writer: &mut MarkdownWriter) -> HandlerOutcome {
192        HandlerOutcome::NoOp
193    }
194}