html_translation_lib/pipeline/

collector.rs

//! 文本收集器模块
//!
//! 从HTML DOM中收集可翻译的文本内容

use crate::error::{TranslationError, TranslationResult};
use markup5ever_rcdom::{Handle, NodeData, RcDom};
use std::collections::HashSet;

/// 文本收集器
pub struct TextCollector {
    /// 跳过的标签集合
    skip_tags: HashSet<String>,
    
    /// 收集属性文本的标签映射
    attribute_tags: std::collections::HashMap<String, Vec<String>>,
}

impl TextCollector {
    /// 创建新的文本收集器
    pub fn new() -> Self {
        let mut skip_tags = HashSet::new();
        skip_tags.insert("script".to_string());
        skip_tags.insert("style".to_string());
        skip_tags.insert("noscript".to_string());
        skip_tags.insert("code".to_string());
        skip_tags.insert("pre".to_string());
        
        let mut attribute_tags = std::collections::HashMap::new();
        attribute_tags.insert("img".to_string(), vec!["alt".to_string(), "title".to_string()]);
        attribute_tags.insert("input".to_string(), vec!["placeholder".to_string(), "value".to_string()]);
        attribute_tags.insert("textarea".to_string(), vec!["placeholder".to_string()]);
        attribute_tags.insert("a".to_string(), vec!["title".to_string()]);
        
        Self {
            skip_tags,
            attribute_tags,
        }
    }
    
    /// 从DOM中收集文本
    pub fn collect_from_dom(&self, dom: &RcDom) -> TranslationResult<Vec<TextItem>> {
        let mut items = Vec::new();
        self.walk_node(&dom.document, &mut items, 0)?;
        Ok(items)
    }
    
    /// 递归遍历DOM节点
    fn walk_node(&self, node: &Handle, items: &mut Vec<TextItem>, depth: usize) -> TranslationResult<()> {
        if depth > 100 {
            return Err(TranslationError::InternalError("DOM深度超限".to_string()));
        }
        
        match &node.data {
            NodeData::Element { name, attrs, .. } => {
                let tag_name = name.local.to_string();
                
                // 检查是否应该跳过此标签
                if !self.skip_tags.contains(&tag_name) {
                    // 收集属性文本
                    if let Some(attr_names) = self.attribute_tags.get(&tag_name) {
                        let attrs = attrs.borrow();
                        for attr in attrs.iter() {
                            let attr_name = attr.name.local.to_string();
                            if attr_names.contains(&attr_name) {
                                let text = attr.value.to_string();
                                if !text.trim().is_empty() {
                                    items.push(TextItem {
                                        text,
                                        text_type: TextType::Attribute(attr_name),
                                        node: node.clone(),
                                        tag_name: tag_name.clone(),
                                    });
                                }
                            }
                        }
                    }
                    
                    // 递归处理子节点
                    for child in node.children.borrow().iter() {
                        self.walk_node(child, items, depth + 1)?;
                    }
                }
            }
            NodeData::Text { contents } => {
                let text = contents.borrow().to_string();
                let trimmed = text.trim();
                if !trimmed.is_empty() {
                    items.push(TextItem {
                        text: trimmed.to_string(),
                        text_type: TextType::Content,
                        node: node.clone(),
                        tag_name: "text".to_string(),
                    });
                }
            }
            _ => {
                // 处理其他节点类型的子节点
                for child in node.children.borrow().iter() {
                    self.walk_node(child, items, depth + 1)?;
                }
            }
        }
        
        Ok(())
    }
}

impl Default for TextCollector {
    fn default() -> Self {
        Self::new()
    }
}

/// 文本项结构体
#[derive(Debug, Clone)]
pub struct TextItem {
    /// 文本内容
    pub text: String,
    
    /// 文本类型
    pub text_type: TextType,
    
    /// 对应的DOM节点
    pub node: Handle,
    
    /// 标签名称
    pub tag_name: String,
}

impl TextItem {
    /// 应用翻译到DOM节点
    pub fn apply_translation(&self, translation: &str) -> TranslationResult<()> {
        match &self.text_type {
            TextType::Content => {
                if let NodeData::Text { contents } = &self.node.data {
                    let mut contents = contents.borrow_mut();
                    *contents = translation.into();
                }
            }
            TextType::Attribute(attr_name) => {
                if let NodeData::Element { attrs, .. } = &self.node.data {
                    let mut attrs = attrs.borrow_mut();
                    for attr in attrs.iter_mut() {
                        if attr.name.local.as_ref() == attr_name {
                            attr.value = translation.into();
                            break;
                        }
                    }
                }
            }
            _ => {}
        }
        Ok(())
    }
}

/// 文本类型枚举
#[derive(Debug, Clone)]
pub enum TextType {
    /// 文本内容
    Content,
    
    /// 标题文本
    Title,
    
    /// 链接文本
    Link,
    
    /// 按钮文本
    Button,
    
    /// 表单标签
    FormLabel,
    
    /// 图片alt属性
    ImageAlt,
    
    /// 工具提示
    Tooltip,
    
    /// 属性值
    Attribute(String),
}

#[cfg(test)]
mod tests {
    use super::*;
    use html5ever::parse_document;
    use html5ever::tendril::TendrilSink;
    use markup5ever_rcdom::RcDom;

    fn parse_html(html: &str) -> RcDom {
        parse_document(RcDom::default(), Default::default())
            .from_utf8()
            .read_from(&mut html.as_bytes())
            .unwrap()
    }

    #[test]
    fn test_text_collector_new() {
        let collector = TextCollector::new();
        
        assert!(collector.skip_tags.contains("script"));
        assert!(collector.skip_tags.contains("style"));
        assert!(collector.skip_tags.contains("noscript"));
        assert!(collector.skip_tags.contains("code"));
        assert!(collector.skip_tags.contains("pre"));
        
        assert!(collector.attribute_tags.contains_key("img"));
        assert!(collector.attribute_tags.contains_key("input"));
        assert!(collector.attribute_tags.contains_key("textarea"));
        assert!(collector.attribute_tags.contains_key("a"));
    }
}