spider_utils/

lib.rs

use hashbrown::{hash_map::Entry, HashMap};
use lazy_static::lazy_static;
use log::{self, warn};
use scraper::{ElementRef, Html, Selector};
use std::{fmt::Debug, hash::Hash};
use sxd_document::parser;
use sxd_xpath::evaluate_xpath;
use tokio_stream::StreamExt;

/// The type of selectors that can be used to query.
#[derive(Default, Debug, Clone)]
pub struct DocumentSelectors<K> {
    /// CSS Selectors.
    pub css: HashMap<K, Vec<Selector>>,
    /// XPath Selectors.
    pub xpath: HashMap<K, Vec<String>>,
}

/// Extracted content from CSS query selectors.
type CSSQueryMap = HashMap<String, Vec<String>>;

lazy_static! {
    /// Xpath factory.
    static ref XPATH_FACTORY: sxd_xpath::Factory = sxd_xpath::Factory::new();
}

/// Check if a selector is a valid xpath
fn is_valid_xpath(expression: &str) -> bool {
    match XPATH_FACTORY.build(expression) {
        Ok(Some(_)) => true,
        Ok(None) => false,
        Err(_) => false,
    }
}

/// Async stream CSS query selector map.
pub async fn css_query_select_map_streamed<K>(
    html: &str,
    selectors: &DocumentSelectors<K>,
) -> CSSQueryMap
where
    K: AsRef<str> + Eq + Hash + Sized,
{
    let mut map: CSSQueryMap = HashMap::new();

    if !selectors.css.is_empty() {
        let mut stream = tokio_stream::iter(&selectors.css);
        let fragment = Box::new(Html::parse_document(html));

        while let Some(selector) = stream.next().await {
            for s in selector.1 {
                for element in fragment.select(s) {
                    process_selector::<K>(element, selector.0, &mut map);
                }
            }
        }
    }

    if !selectors.xpath.is_empty() {
        if let Ok(package) = parser::parse(html) {
            let document = Box::new(package.as_document());

            for selector in selectors.xpath.iter() {
                for s in selector.1 {
                    if let Ok(value) = evaluate_xpath(&document, s) {
                        let text = value.into_string();

                        if !text.is_empty() {
                            match map.entry(selector.0.as_ref().to_string()) {
                                Entry::Occupied(mut entry) => entry.get_mut().push(text),
                                Entry::Vacant(entry) => {
                                    entry.insert(vec![text]);
                                }
                            }
                        }
                    };
                }
            }
        };
    }

    for items in map.values_mut() {
        items.dedup();
    }

    map
}

/// Sync CSS query selector map.
pub fn css_query_select_map<K>(html: &str, selectors: &DocumentSelectors<K>) -> CSSQueryMap
where
    K: AsRef<str> + Eq + Hash + Sized,
{
    let mut map: CSSQueryMap = HashMap::new();

    if !selectors.css.is_empty() {
        let fragment = Box::new(Html::parse_document(html));

        for selector in selectors.css.iter() {
            for s in selector.1 {
                for element in fragment.select(s) {
                    process_selector::<K>(element, selector.0, &mut map);
                }
            }
        }
    }

    if !selectors.xpath.is_empty() {
        if let Ok(package) = parser::parse(html) {
            let document = package.as_document();

            for selector in selectors.xpath.iter() {
                for s in selector.1 {
                    if let Ok(value) = evaluate_xpath(&document, s) {
                        let text = value.into_string();

                        if !text.is_empty() {
                            match map.entry(selector.0.as_ref().to_string()) {
                                Entry::Occupied(mut entry) => entry.get_mut().push(text),
                                Entry::Vacant(entry) => {
                                    entry.insert(vec![text]);
                                }
                            }
                        }
                    };
                }
            }
        };
    }

    map
}

/// Process a single element and update the map with the results.
fn process_selector<K>(element: ElementRef, name: &K, map: &mut CSSQueryMap)
where
    K: AsRef<str> + Eq + Hash + Sized,
{
    let name = name.as_ref();
    let element_name = element.value().name();

    let text = if element_name == "meta" {
        element.attr("content").unwrap_or_default().into()
    } else if element_name == "link" || element_name == "script" || element_name == "styles" {
        match element.attr(if element_name == "link" {
            "href"
        } else {
            "src"
        }) {
            Some(href) => href.into(),
            _ => clean_element_text(&element),
        }
    } else if element_name == "img" || element_name == "source" {
        let mut img_text = String::new();

        if let Some(src) = element.attr("src") {
            if !src.is_empty() {
                img_text.push('[');
                img_text.push_str(src.trim());
                img_text.push(']');
            }
        }
        if let Some(alt) = element.attr("alt") {
            if !alt.is_empty() {
                if img_text.is_empty() {
                    img_text.push_str(alt);
                } else {
                    img_text.push('(');
                    img_text.push('"');
                    img_text.push_str(alt);
                    img_text.push('"');
                    img_text.push(')');
                }
            }
        }

        img_text
    } else {
        clean_element_text(&element)
    };

    if !text.is_empty() {
        match map.entry(name.to_string()) {
            Entry::Occupied(mut entry) => entry.get_mut().push(text),
            Entry::Vacant(entry) => {
                entry.insert(vec![text]);
            }
        }
    }
}

/// get the text extracted.
pub fn clean_element_text(element: &ElementRef) -> String {
    element.text().collect::<Vec<_>>().join(" ")
}

/// Build valid css selectors for extracting. The hashmap takes items with the key for the object key and the value is the css selector.
pub fn build_selectors_base<K, V, S>(selectors: HashMap<K, S>) -> DocumentSelectors<K>
where
    K: AsRef<str> + Eq + Hash + Clone + Debug,
    V: AsRef<str> + Debug + AsRef<str>,
    S: IntoIterator<Item = V>,
{
    let mut valid_selectors: HashMap<K, Vec<Selector>> = HashMap::new();
    let mut valid_selectors_xpath: HashMap<K, Vec<String>> = HashMap::new();

    for (key, selector_set) in selectors {
        let mut selectors_vec = Vec::new();
        let mut selectors_vec_xpath = Vec::new();

        for selector_str in selector_set {
            match Selector::parse(selector_str.as_ref()) {
                Ok(selector) => selectors_vec.push(selector),
                Err(err) => {
                    if is_valid_xpath(selector_str.as_ref()) {
                        selectors_vec_xpath.push(selector_str.as_ref().to_string())
                    } else {
                        warn!(
                            "Failed to parse selector '{}': {:?}",
                            selector_str.as_ref(),
                            err
                        )
                    }
                }
            }
        }

        let has_css_selectors = !selectors_vec.is_empty();
        let has_xpath_selectors = !selectors_vec_xpath.is_empty();

        if has_css_selectors && !has_xpath_selectors {
            valid_selectors.insert(key, selectors_vec);
        } else if !has_css_selectors && has_xpath_selectors {
            valid_selectors_xpath.insert(key, selectors_vec_xpath);
        } else {
            if has_css_selectors {
                valid_selectors.insert(key.clone(), selectors_vec);
            }
            if has_xpath_selectors {
                valid_selectors_xpath.insert(key, selectors_vec_xpath);
            }
        }
    }

    DocumentSelectors {
        css: valid_selectors,
        xpath: valid_selectors_xpath,
    }
}

/// Build valid css selectors for extracting. The hashmap takes items with the key for the object key and the value is the css selector.
#[cfg(not(feature = "indexset"))]
pub fn build_selectors<K, V>(selectors: HashMap<K, hashbrown::HashSet<V>>) -> DocumentSelectors<K>
where
    K: AsRef<str> + Eq + Hash + Clone + Debug,
    V: AsRef<str> + Debug + AsRef<str>,
{
    build_selectors_base::<K, V, hashbrown::HashSet<V>>(selectors)
}

/// Build valid css selectors for extracting. The hashmap takes items with the key for the object key and the value is the css selector.
#[cfg(feature = "indexset")]
pub fn build_selectors<K, V>(selectors: HashMap<K, indexmap::IndexSet<V>>) -> DocumentSelectors<K>
where
    K: AsRef<str> + Eq + Hash + Clone + Debug,
    V: AsRef<str> + Debug + AsRef<str>,
{
    build_selectors_base::<K, V, indexmap::IndexSet<V>>(selectors)
}

#[cfg(not(feature = "indexset"))]
pub type QueryCSSSelectSet<'a> = hashbrown::HashSet<&'a str>;
#[cfg(feature = "indexset")]
pub type QueryCSSSelectSet<'a> = indexmap::IndexSet<&'a str>;
#[cfg(not(feature = "indexset"))]
pub type QueryCSSMap<'a> = HashMap<&'a str, QueryCSSSelectSet<'a>>;
#[cfg(feature = "indexset")]
pub type QueryCSSMap<'a> = HashMap<&'a str, QueryCSSSelectSet<'a>>;

#[cfg(test)]
#[tokio::test]
async fn test_css_query_select_map_streamed() {
    let map = QueryCSSMap::from([("list", QueryCSSSelectSet::from([".list", ".sub-list"]))]);

    let data = css_query_select_map_streamed(
        r#"<html><body><ul class="list"><li>Test</li></ul></body></html>"#,
        &build_selectors(map),
    )
    .await;

    assert!(!data.is_empty(), "CSS extraction failed",);
}

#[test]
fn test_css_query_select_map() {
    let map = QueryCSSMap::from([("list", QueryCSSSelectSet::from([".list", ".sub-list"]))]);
    let data = css_query_select_map(
        r#"<html><body><ul class="list">Test</ul></body></html>"#,
        &build_selectors(map),
    );

    assert!(!data.is_empty(), "CSS extraction failed",);
}

#[cfg(test)]
#[tokio::test]
async fn test_css_query_select_map_streamed_multi_join() {
    let map = QueryCSSMap::from([("list", QueryCSSSelectSet::from([".list", ".sub-list"]))]);
    let data = css_query_select_map_streamed(
        r#"<html>
            <body>
                <ul class="list"><li>First</li></ul>
                <ul class="sub-list"><li>Second</li></ul>
            </body>
        </html>"#,
        &build_selectors(map),
    )
    .await;

    assert!(!data.is_empty(), "CSS extraction failed");
}

#[cfg(test)]
#[tokio::test]
async fn test_xpath_query_select_map_streamed() {
    let map = QueryCSSMap::from([(
        "list",
        QueryCSSSelectSet::from(["//*[@class='list']", "//*[@class='sub-list']"]),
    )]);
    let selectors = build_selectors(map);
    let data = css_query_select_map_streamed(
        r#"<html><body><ul class="list"><li>Test</li></ul></body></html>"#,
        &selectors,
    )
    .await;

    assert!(!data.is_empty(), "Xpath extraction failed",);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_css_query_empty_html() {
        let map = QueryCSSMap::from([("item", QueryCSSSelectSet::from([".item"]))]);
        let data = css_query_select_map("", &build_selectors(map));
        assert!(data.is_empty());
    }

    #[test]
    fn test_css_query_no_matches() {
        let map = QueryCSSMap::from([("item", QueryCSSSelectSet::from([".nonexistent"]))]);
        let data = css_query_select_map(
            r#"<html><body><p class="other">Hello</p></body></html>"#,
            &build_selectors(map),
        );
        assert!(data.is_empty());
    }

    #[test]
    fn test_build_selectors_invalid_css() {
        let map = QueryCSSMap::from([("bad", QueryCSSSelectSet::from(["[[[invalid"]))]);
        let selectors = build_selectors(map);
        // Invalid CSS should be rejected (not panic)
        assert!(selectors.css.is_empty());
    }

    #[test]
    fn test_build_selectors_mixed_css_xpath() {
        let map = QueryCSSMap::from([(
            "mixed",
            QueryCSSSelectSet::from([".valid-css", "//*[@class='xpath']"]),
        )]);
        let selectors = build_selectors(map);
        // Should have CSS selectors and/or XPath selectors
        let has_css = selectors.css.contains_key("mixed");
        let has_xpath = selectors.xpath.contains_key("mixed");
        assert!(has_css || has_xpath);
    }

    #[test]
    fn test_css_query_special_characters() {
        let map = QueryCSSMap::from([("content", QueryCSSSelectSet::from(["p"]))]);
        let data = css_query_select_map(
            r#"<html><body><p>Hello &amp; "world" &lt;test&gt;</p></body></html>"#,
            &build_selectors(map),
        );
        assert!(!data.is_empty());
        let values = data.get("content").unwrap();
        assert!(!values.is_empty());
    }

    #[test]
    fn test_clean_element_text_basic() {
        let html = Html::parse_fragment("<p>Hello <b>World</b></p>");
        let selector = Selector::parse("p").unwrap();
        if let Some(element) = html.select(&selector).next() {
            let text = clean_element_text(&element);
            assert!(text.contains("Hello"));
            assert!(text.contains("World"));
        }
    }

    #[test]
    fn test_process_selector_img_element() {
        let html = Html::parse_fragment(r#"<img src="photo.jpg" alt="A photo">"#);
        let selector = Selector::parse("img").unwrap();
        let mut map: HashMap<String, Vec<String>> = HashMap::new();

        if let Some(element) = html.select(&selector).next() {
            process_selector::<&str>(element, &"image", &mut map);
        }
        assert!(map.contains_key("image"));
        let vals = &map["image"];
        assert!(!vals.is_empty());
        // Should contain src in brackets and alt in parens
        assert!(vals[0].contains("photo.jpg"));
    }

    #[test]
    fn test_process_selector_meta_element() {
        let html = Html::parse_document(
            r#"<html><head><meta name="description" content="Test description"></head><body></body></html>"#,
        );
        let selector = Selector::parse("meta[name='description']").unwrap();
        let mut map: HashMap<String, Vec<String>> = HashMap::new();

        if let Some(element) = html.select(&selector).next() {
            process_selector::<&str>(element, &"desc", &mut map);
        }
        assert!(map.contains_key("desc"));
        assert_eq!(map["desc"][0], "Test description");
    }
}