Implement Html::select

indefini · Jan 15, 2016 · 49e85a9 · 49e85a9
1 parent c744c34
commit 49e85a9
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 1 deletion.
diff --git a/examples/select.rs b/examples/select.rs
@@ -0,0 +1,34 @@
+extern crate scraper;
+
+use std::io;
+
+use scraper::{Html, Selector};
+
+static HTML: &'static str = r##"
+<!DOCTYPE html>
+<title>Hello, world!</title>
+<meta charset="utf-8">
+<header id="header">
+  <h1>Title</h1>
+  <p class="tagline">Tagline</p>
+  <nav>
+    <ul>
+      <li><a href="#">Nav Link</a></li>
+    </ul>
+  </nav>
+</header>
+<main>
+  <p>Content</p>
+</main>
+"##;
+
+fn main() {
+    let html = Html::parse_document(HTML);
+    println!("{:#?}", html);
+    let mut input = String::new();
+    io::stdin().read_line(&mut input).unwrap();
+    let selector = Selector::parse(&input).unwrap();
+    for node in html.select(&selector) {
+        println!("{:?}", node.value());
+    }
+}
diff --git a/src/html/mod.rs b/src/html/mod.rs
@@ -3,11 +3,14 @@
 use std::borrow::Cow;
 
 use ego_tree::Tree;
+use ego_tree::iter::Nodes;
 use html5ever::driver;
 use html5ever::tree_builder::QuirksMode;
 use tendril::StrTendril;
 
 use node::Node;
+use node_ref::NodeRef;
+use selector::Selector;
 
 /// An HTML tree.
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -60,6 +63,35 @@ impl Html {
             Default::default()
         )
     }
+
+    /// Returns an iterator over elements matching a selector.
+    pub fn select<'a, 'b>(&'a self, selector: &'b Selector) -> Select<'a, 'b> {
+        Select {
+            selector: selector,
+            inner: self.tree.nodes(),
+        }
+    }
+}
+
+/// Iterator over elements matching a selector.
+#[derive(Debug)]
+pub struct Select<'a, 'b> {
+    inner: Nodes<'a, Node>,
+    selector: &'b Selector,
+}
+
+impl<'a, 'b> Iterator for Select<'a, 'b> {
+    type Item = NodeRef<'a>;
+
+    fn next(&mut self) -> Option<NodeRef<'a>> {
+        for node in self.inner.by_ref() {
+            let node_ref = NodeRef(node);
+            if node.value().is_element() && self.selector.matches(node_ref) {
+                return Some(node_ref);
+            }
+        }
+        None
+    }
 }
 
 mod tree_sink;
diff --git a/src/node_ref/mod.rs b/src/node_ref/mod.rs
@@ -7,7 +7,7 @@ use ego_tree;
 use node::Node;
 
 /// Wrapper around a reference to an HTML node.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct NodeRef<'a>(pub ego_tree::NodeRef<'a, Node>);
 
 impl<'a> Deref for NodeRef<'a> {

diff --git a/src/selector.rs b/src/selector.rs
@@ -1,8 +1,11 @@
 //! CSS selectors.
 
 use cssparser::Parser;
+use selectors::matching;
 use selectors::parser::{self, ParserContext};
 
+use node_ref::NodeRef;
+
 /// Wrapper around CSS selectors.
 #[derive(Debug, Clone, PartialEq)]
 pub struct Selector {
@@ -18,4 +21,13 @@ impl Selector {
         let selectors = try!(parser::parse_selector_list(&context, &mut parser));
         Ok(Selector { selectors: selectors })
     }
+
+    /// Returns true if the referenced element matches this selector.
+    ///
+    /// # Panics
+    ///
+    /// Panics if referenced node is not an element.
+    pub fn matches(&self, node: NodeRef) -> bool {
+        matching::matches(&self.selectors, &node, None)
+    }
 }