use ego_tree::NodeId;
use scraper::{Html, Selector};
use std::sync::LazyLock;
pub static BODY_SELECTOR: LazyLock<Selector> = LazyLock::new(|| {
Selector::parse("body").expect("Can't be (parsing body selector)")
});
#[derive(Debug, Clone, Default)]
pub struct NodeMetrics {
pub char_count: u32,
pub tag_count: u32,
pub link_char_count: u32,
pub link_tag_count: u32,
}
pub trait TreeBuilder {
fn build_metrics(&self, node_id: NodeId) -> NodeMetrics;
fn get_children(&self, node_id: NodeId) -> Vec<NodeId>;
fn get_parent(&self, node_id: NodeId) -> Option<NodeId>;
}
pub struct HtmlTreeBuilder<'a> {
document: &'a Html,
}
impl<'a> HtmlTreeBuilder<'a> {
pub fn new(document: &'a Html) -> Self {
Self { document }
}
}
impl TreeBuilder for HtmlTreeBuilder<'_> {
fn build_metrics(&self, node_id: NodeId) -> NodeMetrics {
let node = self.document.tree.get(node_id).unwrap();
let mut metrics = NodeMetrics {
char_count: 0,
tag_count: 0,
link_char_count: 0,
link_tag_count: 0,
};
match node.value() {
scraper::Node::Text(text) => {
let clean_text = text.trim();
if !crate::utils::is_non_content_text(clean_text) {
metrics.char_count =
crate::unicode::count_graphemes(clean_text);
}
}
scraper::Node::Element(elem) => {
metrics.tag_count = 1;
if elem.name() == "a"
|| elem.name() == "button"
|| elem.name() == "select"
{
metrics.link_tag_count = 1;
}
}
_ => {}
}
metrics
}
fn get_children(&self, node_id: NodeId) -> Vec<NodeId> {
self.document
.tree
.get(node_id)
.map(|node| {
node.children()
.filter(|child| match child.value() {
scraper::Node::Element(elem) => {
!crate::utils::should_skip_element(elem)
}
scraper::Node::Text(text) => {
!crate::utils::is_non_content_text(text)
}
scraper::Node::Comment(_) | scraper::Node::Document => {
false
}
_ => true,
})
.map(|child| child.id())
.collect()
})
.unwrap_or_default()
}
fn get_parent(&self, node_id: NodeId) -> Option<NodeId> {
self.document
.tree
.get(node_id)
.and_then(|node| node.parent())
.map(|parent| parent.id())
}
}
impl NodeMetrics {
pub fn new() -> Self {
Self::default()
}
pub fn combine(&mut self, other: &NodeMetrics) {
self.char_count += other.char_count;
self.tag_count += other.tag_count;
self.link_char_count += other.link_char_count;
self.link_tag_count += other.link_tag_count;
}
pub fn calculate_simple_density(&self) -> f32 {
if self.tag_count == 0 {
0.0
} else {
self.char_count as f32 / self.tag_count as f32
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use scraper::Html;
const TEST_HTML: &str = r#"
<html>
<body>
<div class="content">
Some text here
<a href="\#">A link</a>
<p>More content</p>
<button>Click me</button>
<script>console.log('skip');</script>
<style>.skip{}</style>
</div>
<div class="sidebar">
<select>
<option>Option 1</option>
</select>
</div>
</body>
</html>
"#;
#[test]
fn test_body_selector_initialization() {
let _ = &*BODY_SELECTOR;
}
#[test]
fn test_node_metrics() {
let document = Html::parse_document(TEST_HTML);
let builder = HtmlTreeBuilder::new(&document);
let content_div = document
.select(&scraper::Selector::parse("div.content").unwrap())
.next()
.unwrap();
let metrics = builder.build_metrics(content_div.id());
assert_eq!(metrics.char_count, 0); assert_eq!(metrics.tag_count, 1); assert_eq!(metrics.link_tag_count, 0); }
#[test]
fn test_get_children_filters() {
let document = Html::parse_document(TEST_HTML);
let builder = HtmlTreeBuilder::new(&document);
let body = document
.select(&scraper::Selector::parse("body").unwrap())
.next()
.unwrap();
let children = builder.get_children(body.id());
assert_eq!(children.len(), 2);
}
#[test]
fn test_link_metrics() {
let document = Html::parse_document(TEST_HTML);
let builder = HtmlTreeBuilder::new(&document);
let link = document
.select(&scraper::Selector::parse("a").unwrap())
.next()
.unwrap();
let metrics = builder.build_metrics(link.id());
assert_eq!(metrics.link_tag_count, 1);
let button = document
.select(&scraper::Selector::parse("button").unwrap())
.next()
.unwrap();
let metrics = builder.build_metrics(button.id());
assert_eq!(metrics.link_tag_count, 1);
let select = document
.select(&scraper::Selector::parse("select").unwrap())
.next()
.unwrap();
let metrics = builder.build_metrics(select.id());
assert_eq!(metrics.link_tag_count, 1);
}
#[test]
fn test_text_metrics() {
let document = Html::parse_document(TEST_HTML);
let builder = HtmlTreeBuilder::new(&document);
let text_node = document
.select(&scraper::Selector::parse(".content").unwrap())
.next()
.unwrap()
.first_child()
.unwrap();
let metrics = builder.build_metrics(text_node.id());
assert_eq!(metrics.char_count, 14); assert_eq!(metrics.tag_count, 0);
assert_eq!(metrics.link_tag_count, 0);
}
#[test]
fn test_get_parent() {
let document = Html::parse_document(TEST_HTML);
let builder = HtmlTreeBuilder::new(&document);
let content_div = document
.select(&scraper::Selector::parse("div.content").unwrap())
.next()
.unwrap();
let parent = builder.get_parent(content_div.id());
assert!(parent.is_some());
let body = document
.select(&scraper::Selector::parse("body").unwrap())
.next()
.unwrap();
assert_eq!(parent.unwrap(), body.id());
}
}