From 325fbe47d1545b720584208d2b9d44d16633d781 Mon Sep 17 00:00:00 2001
From: Wouter Scherphof <wouter.scherphof@gmail.com>
Date: Wed, 20 Mar 2013 09:38:43 +0100
Subject: [PATCH] closes #8

---
 ElementNode.lua | 28 +++++++++++++++++++++++-----
 test.lua        | 21 +++++++++++----------
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/ElementNode.lua b/ElementNode.lua
index 9b23b36..fdeebc4 100644
--- a/ElementNode.lua
+++ b/ElementNode.lua
@@ -17,7 +17,9 @@ function ElementNode:new(nameortext, node, descend, openstart, openend)
     deepernodes = Set:new(),
     deeperelements = {}, deeperattributes = {}, deeperids = {}, deeperclasses = {}
   }
-  if not node then
+  if nameortext == "container" then
+    instance.root = node
+  elseif not node then
     instance.name = "root"
     instance.root = instance
     instance._text = nameortext
@@ -60,9 +62,7 @@ function ElementNode:addattribute(k, v)
 end
 
 local function insert(list, name, node)
-  if not list[name] then
-    list[name] = Set:new()
-  end
+  list[name] = list[name] or Set:new()
   list[name]:add(node)
 end
 
@@ -113,7 +113,25 @@ local function select(self, s)
     subjects = Set:new(resultset)
     ::nextpart::
   end
-  return resultset:tolist()
+  -- construct a container node for the resultset, so that we can :select() on it
+  local ret = ElementNode:new("container", self)
+  for node in pairs(resultset) do
+    table.insert(ret.nodes, node)
+    ret.deepernodes = ret.deepernodes + node.deepernodes
+    for listname,list in pairs({
+      deeperelements = node.deeperelements,
+      deeperattributes = node.deeperattributes,
+      deeperids = node.deeperids,
+      deeperclasses = node.deeperclasses
+    }) do
+      local target = ret[listname]
+      for k,set in pairs(list) do
+        -- Set.__add will create an empty Set if not target[k]
+        target[k] = target[k] + set
+      end
+    end
+  end
+  return ret
 end
 
 function ElementNode:select(s) return select(self, s) end
diff --git a/test.lua b/test.lua
index 7774288..9d3ffc9 100644
--- a/test.lua
+++ b/test.lua
@@ -27,10 +27,10 @@ local function select( s )
   print ""
   print("->", s)
   local tags = root:select(s)
-  for i,t in ipairs(tags) do
+  for i,t in ipairs(tags.nodes) do
     print(t.name)
   end
-  print(# tags)
+  print(# tags.nodes)
 end
 select("*")
 select("link")
@@ -52,7 +52,7 @@ select("body [class]")
 select("body > [class]")
 
 local sel, chapters = root("ol.chapters > li"), {}
-for i,v in ipairs(sel) do
+for _,v in ipairs(sel.nodes) do
   table.insert(chapters, v:getcontent())
 end
 print("\nchapters")
@@ -60,13 +60,11 @@ for i,v in ipairs(chapters) do
   print(i, v)
 end
 
-local sel, contacts = root("ul.contacts > li"), {}
-for i,v in ipairs(sel) do
-  local c = {}
-  for fi,fv in ipairs(v("span[class]")) do
-    c[fv.classes[1]] = fv:getcontent()
-  end
-  contacts[v.id] = c
+local sel, contacts = root("ul.contacts > li")("span[class]"), {}
+for _,v in ipairs(sel.nodes) do
+  local id = v.parent.parent.id -- li > a > span
+  contacts[id] = contacts[id] or {}
+  contacts[id][v.classes[1]] = v:getcontent()
 end
 print("\ncontacts")
 for k,v in pairs(contacts) do
@@ -75,3 +73,6 @@ for k,v in pairs(contacts) do
     print(fk, fv)
   end
 end
+
+
+