Skip to main content

trek_rs/
extractor.rs

1//! Site-specific content extractors
2
3use crate::types::ExtractedContent;
4use eyre::Result;
5use serde_json::Value;
6use tracing::{debug, instrument};
7
8/// Trait for site-specific extractors
9pub trait Extractor: Send + Sync {
10    /// Check if this extractor can handle the current document
11    fn can_extract(&self, url: &str, schema_org_data: &[Value]) -> bool;
12
13    /// Extract content from HTML string
14    fn extract_from_html(&self, html: &str) -> Result<ExtractedContent>;
15
16    /// Get the name of this extractor
17    fn name(&self) -> &'static str;
18}
19
20/// Registry for site-specific extractors
21pub struct ExtractorRegistry {
22    extractors: Vec<Box<dyn Extractor>>,
23}
24
25impl std::fmt::Debug for ExtractorRegistry {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        f.debug_struct("ExtractorRegistry")
28            .field("extractors_count", &self.extractors.len())
29            .finish()
30    }
31}
32
33impl Default for ExtractorRegistry {
34    fn default() -> Self {
35        Self::new()
36    }
37}
38
39impl ExtractorRegistry {
40    /// Create a new extractor registry
41    pub fn new() -> Self {
42        // Register built-in extractors
43        // TODO: Add extractors as we implement them
44
45        Self {
46            extractors: Vec::new(),
47        }
48    }
49
50    /// Register a new extractor
51    pub fn register(&mut self, extractor: Box<dyn Extractor>) {
52        debug!("Registering extractor: {}", extractor.name());
53        self.extractors.push(extractor);
54    }
55
56    /// Find an extractor that can handle the current document
57    #[instrument(skip(self, schema_org_data))]
58    pub fn find_extractor_from_data(
59        &self,
60        url: &str,
61        schema_org_data: &[Value],
62    ) -> Option<&dyn Extractor> {
63        for extractor in &self.extractors {
64            if extractor.can_extract(url, schema_org_data) {
65                debug!("Found matching extractor: {}", extractor.name());
66                return Some(extractor.as_ref());
67            }
68        }
69        None
70    }
71}
72
73/// Generic content extractor (fallback)
74pub struct GenericExtractor;
75
76impl Extractor for GenericExtractor {
77    fn can_extract(&self, _url: &str, _schema_org_data: &[Value]) -> bool {
78        // Generic extractor should not be used as a site-specific extractor
79        // It's better to fall back to the main extraction logic
80        false
81    }
82
83    fn extract_from_html(&self, html: &str) -> Result<ExtractedContent> {
84        // Basic extraction logic
85        let mut content = ExtractedContent::default();
86
87        // Extract title from HTML
88        if let Some(title_start) = html.find("<title>") {
89            if let Some(title_end) = html[title_start..].find("</title>") {
90                let title = &html[title_start + 7..title_start + title_end];
91                content.title = Some(title.trim().to_string());
92            }
93        }
94
95        Ok(content)
96    }
97
98    fn name(&self) -> &'static str {
99        "generic"
100    }
101}
102
103#[cfg(test)]
104#[allow(clippy::disallowed_methods)] // OK to use unwrap/expect in tests
105mod tests {
106    use super::*;
107
108    #[test]
109    fn test_generic_extractor() {
110        let extractor = GenericExtractor;
111        let html = r"<html><head><title>Test Title</title></head></html>";
112
113        let result = extractor.extract_from_html(html).unwrap();
114        assert_eq!(result.title, Some("Test Title".to_string()));
115    }
116
117    struct TestExtractor;
118
119    impl Extractor for TestExtractor {
120        fn can_extract(&self, url: &str, _schema_org_data: &[Value]) -> bool {
121            url.contains("test.com")
122        }
123        fn extract_from_html(&self, _html: &str) -> Result<ExtractedContent> {
124            Ok(ExtractedContent::default())
125        }
126        fn name(&self) -> &'static str {
127            "test"
128        }
129    }
130
131    #[test]
132    fn test_registry() {
133        let mut registry = ExtractorRegistry::new();
134        registry.register(Box::new(GenericExtractor));
135
136        // GenericExtractor should not match any URL (it returns false for can_extract)
137        let extractor = registry.find_extractor_from_data("https://example.com", &[]);
138        assert!(extractor.is_none());
139
140        // Test that registry can find extractors when they match
141        registry.register(Box::new(TestExtractor));
142        let extractor = registry.find_extractor_from_data("https://test.com", &[]);
143        assert!(extractor.is_some());
144        assert_eq!(extractor.unwrap().name(), "test");
145    }
146}