dom_content_extraction/
tree.rs

1use ego_tree::NodeId;
2use scraper::{Html, Selector};
3use std::sync::LazyLock;
4
5/// Selector for `<body>` tag
6pub static BODY_SELECTOR: LazyLock<Selector> = LazyLock::new(|| {
7    Selector::parse("body").expect("Can't be (parsing body selector)")
8});
9
10#[derive(Debug, Clone, Default)]
11pub struct NodeMetrics {
12    pub char_count: u32,
13    pub tag_count: u32,
14    pub link_char_count: u32,
15    pub link_tag_count: u32,
16}
17
18pub trait TreeBuilder {
19    fn build_metrics(&self, node_id: NodeId) -> NodeMetrics;
20    fn get_children(&self, node_id: NodeId) -> Vec<NodeId>;
21    fn get_parent(&self, node_id: NodeId) -> Option<NodeId>;
22}
23
24pub struct HtmlTreeBuilder<'a> {
25    document: &'a Html,
26}
27
28impl<'a> HtmlTreeBuilder<'a> {
29    pub fn new(document: &'a Html) -> Self {
30        Self { document }
31    }
32}
33
34impl TreeBuilder for HtmlTreeBuilder<'_> {
35    fn build_metrics(&self, node_id: NodeId) -> NodeMetrics {
36        let node = self.document.tree.get(node_id).unwrap();
37
38        let mut metrics = NodeMetrics {
39            char_count: 0,
40            tag_count: 0,
41            link_char_count: 0,
42            link_tag_count: 0,
43        };
44
45        match node.value() {
46            scraper::Node::Text(text) => {
47                // NOTE: old method calculation
48                // metrics.char_count = text.trim().len() as u32;
49                let clean_text = text.trim();
50                if !crate::utils::is_non_content_text(clean_text) {
51                    metrics.char_count =
52                        crate::unicode::count_graphemes(clean_text);
53                }
54            }
55            scraper::Node::Element(elem) => {
56                metrics.tag_count = 1;
57                if elem.name() == "a"
58                    || elem.name() == "button"
59                    || elem.name() == "select"
60                {
61                    metrics.link_tag_count = 1;
62                }
63            }
64            _ => {}
65        }
66
67        metrics
68    }
69
70    fn get_children(&self, node_id: NodeId) -> Vec<NodeId> {
71        self.document
72            .tree
73            .get(node_id)
74            .map(|node| {
75                node.children()
76                    .filter(|child| match child.value() {
77                        scraper::Node::Element(elem) => {
78                            !crate::utils::should_skip_element(elem)
79                        }
80                        scraper::Node::Text(text) => {
81                            !crate::utils::is_non_content_text(text)
82                        }
83                        scraper::Node::Comment(_) | scraper::Node::Document => {
84                            false
85                        }
86                        _ => true,
87                    })
88                    .map(|child| child.id())
89                    .collect()
90            })
91            .unwrap_or_default()
92    }
93
94    fn get_parent(&self, node_id: NodeId) -> Option<NodeId> {
95        self.document
96            .tree
97            .get(node_id)
98            .and_then(|node| node.parent())
99            .map(|parent| parent.id())
100    }
101}
102
103impl NodeMetrics {
104    pub fn new() -> Self {
105        Self::default()
106    }
107
108    // Add metrics from another (child) node
109    pub fn combine(&mut self, other: &NodeMetrics) {
110        self.char_count += other.char_count;
111        self.tag_count += other.tag_count;
112        self.link_char_count += other.link_char_count;
113        self.link_tag_count += other.link_tag_count;
114    }
115
116    // Calculate simple density (chars/tags ratio)
117    pub fn calculate_simple_density(&self) -> f32 {
118        if self.tag_count == 0 {
119            0.0
120        } else {
121            self.char_count as f32 / self.tag_count as f32
122        }
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use scraper::Html;
130
131    const TEST_HTML: &str = r#"
132        <html>
133        <body>
134            <div class="content">
135                Some text here
136                <a href="\#">A link</a>
137                <p>More content</p>
138                <button>Click me</button>
139                <script>console.log('skip');</script>
140                <style>.skip{}</style>
141            </div>
142            <div class="sidebar">
143                <select>
144                    <option>Option 1</option>
145                </select>
146            </div>
147        </body>
148        </html>
149    "#;
150
151    #[test]
152    fn test_body_selector_initialization() {
153        // This will force the LazyLock to initialize
154        let _ = &*BODY_SELECTOR;
155    }
156
157    #[test]
158    fn test_node_metrics() {
159        let document = Html::parse_document(TEST_HTML);
160        let builder = HtmlTreeBuilder::new(&document);
161
162        // Get content div
163        let content_div = document
164            .select(&scraper::Selector::parse("div.content").unwrap())
165            .next()
166            .unwrap();
167
168        let metrics = builder.build_metrics(content_div.id());
169        // notice that we do not count children nodes metrics
170        assert_eq!(metrics.char_count, 0); // div itself has no direct text
171        assert_eq!(metrics.tag_count, 1); // div counts as one tag
172        assert_eq!(metrics.link_tag_count, 0); // div is not a link
173    }
174
175    #[test]
176    fn test_get_children_filters() {
177        let document = Html::parse_document(TEST_HTML);
178        let builder = HtmlTreeBuilder::new(&document);
179
180        let body = document
181            .select(&scraper::Selector::parse("body").unwrap())
182            .next()
183            .unwrap();
184
185        let children = builder.get_children(body.id());
186
187        // Should get both divs but skip script and style
188        assert_eq!(children.len(), 2);
189    }
190
191    #[test]
192    fn test_link_metrics() {
193        let document = Html::parse_document(TEST_HTML);
194        let builder = HtmlTreeBuilder::new(&document);
195
196        // Test link element
197        let link = document
198            .select(&scraper::Selector::parse("a").unwrap())
199            .next()
200            .unwrap();
201
202        let metrics = builder.build_metrics(link.id());
203        assert_eq!(metrics.link_tag_count, 1);
204
205        // Test button
206        let button = document
207            .select(&scraper::Selector::parse("button").unwrap())
208            .next()
209            .unwrap();
210
211        let metrics = builder.build_metrics(button.id());
212        assert_eq!(metrics.link_tag_count, 1);
213
214        // Test select
215        let select = document
216            .select(&scraper::Selector::parse("select").unwrap())
217            .next()
218            .unwrap();
219
220        let metrics = builder.build_metrics(select.id());
221        assert_eq!(metrics.link_tag_count, 1);
222    }
223
224    #[test]
225    fn test_text_metrics() {
226        let document = Html::parse_document(TEST_HTML);
227        let builder = HtmlTreeBuilder::new(&document);
228
229        // Find text node
230        let text_node = document
231            .select(&scraper::Selector::parse(".content").unwrap())
232            .next()
233            .unwrap()
234            .first_child()
235            .unwrap();
236
237        let metrics = builder.build_metrics(text_node.id());
238        assert_eq!(metrics.char_count, 14); // "Some text here"
239        assert_eq!(metrics.tag_count, 0);
240        assert_eq!(metrics.link_tag_count, 0);
241    }
242
243    #[test]
244    fn test_get_parent() {
245        let document = Html::parse_document(TEST_HTML);
246        let builder = HtmlTreeBuilder::new(&document);
247
248        let content_div = document
249            .select(&scraper::Selector::parse("div.content").unwrap())
250            .next()
251            .unwrap();
252
253        let parent = builder.get_parent(content_div.id());
254        assert!(parent.is_some());
255
256        // Body should be parent
257        let body = document
258            .select(&scraper::Selector::parse("body").unwrap())
259            .next()
260            .unwrap();
261
262        assert_eq!(parent.unwrap(), body.id());
263    }
264}
dom_content_extraction/tree.rs

dom_content_extraction/
tree.rs
