browser_use/dom/
selector_map.rs

1use indexmap::IndexMap;
2use serde::{Deserialize, Serialize};
3
4/// Information needed to locate an element
5#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
6pub struct ElementSelector {
7    /// CSS selector for the element
8    pub css_selector: String,
9
10    /// XPath selector (alternative to CSS)
11    #[serde(skip_serializing_if = "Option::is_none")]
12    pub xpath: Option<String>,
13
14    /// Element's tag name
15    pub tag_name: String,
16
17    /// Element's ID attribute (if any)
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub id: Option<String>,
20
21    /// Element's text content (truncated for display)
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub text: Option<String>,
24}
25
26impl ElementSelector {
27    /// Create a new ElementSelector with CSS selector
28    pub fn new(css_selector: impl Into<String>, tag_name: impl Into<String>) -> Self {
29        Self {
30            css_selector: css_selector.into(),
31            xpath: None,
32            tag_name: tag_name.into(),
33            id: None,
34            text: None,
35        }
36    }
37
38    /// Builder method: set XPath
39    pub fn with_xpath(mut self, xpath: impl Into<String>) -> Self {
40        self.xpath = Some(xpath.into());
41        self
42    }
43
44    /// Builder method: set ID
45    pub fn with_id(mut self, id: impl Into<String>) -> Self {
46        self.id = Some(id.into());
47        self
48    }
49
50    /// Builder method: set text content
51    pub fn with_text(mut self, text: impl Into<String>) -> Self {
52        self.text = Some(text.into());
53        self
54    }
55
56    /// Get the best selector to use (CSS preferred)
57    pub fn best_selector(&self) -> &str {
58        &self.css_selector
59    }
60}
61
62/// Map of element indices to their selectors
63/// Uses IndexMap to preserve insertion order
64#[derive(Debug, Clone, Default)]
65pub struct SelectorMap {
66    /// Map from index to selector information
67    map: IndexMap<usize, ElementSelector>,
68
69    /// Next available index
70    next_index: usize,
71}
72
73impl SelectorMap {
74    /// Create a new empty SelectorMap
75    pub fn new() -> Self {
76        Self {
77            map: IndexMap::new(),
78            next_index: 0,
79        }
80    }
81
82    /// Register a new element and return its assigned index
83    pub fn register(&mut self, selector: ElementSelector) -> usize {
84        let index = self.next_index;
85        self.map.insert(index, selector);
86        self.next_index += 1;
87        index
88    }
89
90    /// Get selector by index
91    pub fn get(&self, index: usize) -> Option<&ElementSelector> {
92        self.map.get(&index)
93    }
94
95    /// Get mutable selector by index
96    pub fn get_mut(&mut self, index: usize) -> Option<&mut ElementSelector> {
97        self.map.get_mut(&index)
98    }
99
100    /// Check if index exists
101    pub fn contains(&self, index: usize) -> bool {
102        self.map.contains_key(&index)
103    }
104
105    /// Remove an element by index
106    pub fn remove(&mut self, index: usize) -> Option<ElementSelector> {
107        self.map.shift_remove(&index)
108    }
109
110    /// Get the number of registered elements
111    pub fn len(&self) -> usize {
112        self.map.len()
113    }
114
115    /// Check if the map is empty
116    pub fn is_empty(&self) -> bool {
117        self.map.is_empty()
118    }
119
120    /// Clear all elements
121    pub fn clear(&mut self) {
122        self.map.clear();
123        self.next_index = 0;
124    }
125
126    /// Iterate over all (index, selector) pairs
127    pub fn iter(&self) -> impl Iterator<Item = (&usize, &ElementSelector)> {
128        self.map.iter()
129    }
130
131    /// Get all indices
132    pub fn indices(&self) -> impl Iterator<Item = &usize> {
133        self.map.keys()
134    }
135
136    /// Get all selectors
137    pub fn selectors(&self) -> impl Iterator<Item = &ElementSelector> {
138        self.map.values()
139    }
140
141    /// Find index by CSS selector
142    pub fn find_by_css_selector(&self, css_selector: &str) -> Option<usize> {
143        self.map
144            .iter()
145            .find(|(_, sel)| sel.css_selector == css_selector)
146            .map(|(idx, _)| *idx)
147    }
148
149    /// Find index by element ID
150    pub fn find_by_id(&self, id: &str) -> Option<usize> {
151        self.map
152            .iter()
153            .find(|(_, sel)| sel.id.as_deref() == Some(id))
154            .map(|(idx, _)| *idx)
155    }
156
157    /// Export to JSON for debugging
158    pub fn to_json(&self) -> Result<String, serde_json::Error> {
159        serde_json::to_string_pretty(&self.map)
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    #[test]
168    fn test_element_selector() {
169        let selector = ElementSelector::new("#my-button", "button")
170            .with_id("my-button")
171            .with_text("Click me");
172
173        assert_eq!(selector.css_selector, "#my-button");
174        assert_eq!(selector.tag_name, "button");
175        assert_eq!(selector.id, Some("my-button".to_string()));
176        assert_eq!(selector.text, Some("Click me".to_string()));
177        assert_eq!(selector.best_selector(), "#my-button");
178    }
179
180    #[test]
181    fn test_selector_map_register() {
182        let mut map = SelectorMap::new();
183
184        let selector1 = ElementSelector::new("#btn1", "button");
185        let selector2 = ElementSelector::new("#btn2", "button");
186
187        let idx1 = map.register(selector1);
188        let idx2 = map.register(selector2);
189
190        assert_eq!(idx1, 0);
191        assert_eq!(idx2, 1);
192        assert_eq!(map.len(), 2);
193    }
194
195    #[test]
196    fn test_selector_map_get() {
197        let mut map = SelectorMap::new();
198
199        let selector = ElementSelector::new("#test", "div").with_id("test");
200        let index = map.register(selector);
201
202        let retrieved = map.get(index).unwrap();
203        assert_eq!(retrieved.css_selector, "#test");
204        assert_eq!(retrieved.id, Some("test".to_string()));
205    }
206
207    #[test]
208    fn test_selector_map_remove() {
209        let mut map = SelectorMap::new();
210
211        let selector = ElementSelector::new("#remove-me", "span");
212        let index = map.register(selector);
213
214        assert!(map.contains(index));
215
216        let removed = map.remove(index);
217        assert!(removed.is_some());
218        assert!(!map.contains(index));
219        assert_eq!(map.len(), 0);
220    }
221
222    #[test]
223    fn test_selector_map_clear() {
224        let mut map = SelectorMap::new();
225
226        map.register(ElementSelector::new("#one", "div"));
227        map.register(ElementSelector::new("#two", "div"));
228        map.register(ElementSelector::new("#three", "div"));
229
230        assert_eq!(map.len(), 3);
231
232        map.clear();
233
234        assert_eq!(map.len(), 0);
235        assert!(map.is_empty());
236    }
237
238    #[test]
239    fn test_selector_map_find() {
240        let mut map = SelectorMap::new();
241
242        let selector1 = ElementSelector::new("#btn1", "button").with_id("btn1");
243        let selector2 = ElementSelector::new(".link", "a").with_id("link1");
244
245        let idx1 = map.register(selector1);
246        let _idx2 = map.register(selector2);
247
248        let found = map.find_by_css_selector("#btn1");
249        assert_eq!(found, Some(idx1));
250
251        let found_by_id = map.find_by_id("link1");
252        assert_eq!(found_by_id, Some(1));
253
254        let not_found = map.find_by_css_selector("#nonexistent");
255        assert!(not_found.is_none());
256    }
257
258    #[test]
259    fn test_selector_map_iteration() {
260        let mut map = SelectorMap::new();
261
262        map.register(ElementSelector::new("#one", "div"));
263        map.register(ElementSelector::new("#two", "div"));
264        map.register(ElementSelector::new("#three", "div"));
265
266        let indices: Vec<_> = map.indices().copied().collect();
267        assert_eq!(indices, vec![0, 1, 2]);
268
269        let css_selectors: Vec<_> = map.selectors().map(|s| s.css_selector.clone()).collect();
270        assert_eq!(css_selectors, vec!["#one", "#two", "#three"]);
271    }
272
273    #[test]
274    fn test_selector_serialization() {
275        let selector = ElementSelector::new("#test", "button")
276            .with_id("test")
277            .with_text("Test Button");
278
279        let json = serde_json::to_string(&selector).unwrap();
280        let deserialized: ElementSelector = serde_json::from_str(&json).unwrap();
281
282        assert_eq!(selector, deserialized);
283    }
284
285    #[test]
286    fn test_selector_map_to_json() {
287        let mut map = SelectorMap::new();
288
289        map.register(ElementSelector::new("#btn", "button").with_text("Click"));
290        map.register(ElementSelector::new("#link", "a").with_text("Visit"));
291
292        let json = map.to_json().unwrap();
293        assert!(json.contains("#btn"));
294        assert!(json.contains("#link"));
295        assert!(json.contains("Click"));
296        assert!(json.contains("Visit"));
297    }
298}
browser_use/dom/selector_map.rs

browser_use/dom/
selector_map.rs
