find many attrs from one single query

Author: felipejfc

gluahttpscrape.go

@@ -18,12 +18,15 @@ func NewHttpScrapeModule() *httpScrapeModule {
 
 func (h *httpScrapeModule) Loader(L *lua.LState) int {
 	mod := L.SetFuncs(L.NewTable(), map[string]lua.LGFunction{
-		"find_attr_by_class": h.findAttrByClass,
-		"find_text_by_class": h.findTextByClass,
-		"find_attr_by_id":    h.findAttrById,
-		"find_text_by_id":    h.findTextById,
-		"find_attr_by_tag":   h.findAttrByTag,
-		"find_text_by_tag":   h.findTextByTag,
+		"find_attr_by_class":  h.findAttrByClass,
+		"find_attr_by_id":     h.findAttrById,
+		"find_attr_by_tag":    h.findAttrByTag,
+		"find_attrs_by_class": h.findAttrsByClass,
+		"find_attrs_by_id":    h.findAttrsById,
+		"find_attrs_by_tag":   h.findAttrsByTag,
+		"find_text_by_id":     h.findTextById,
+		"find_text_by_class":  h.findTextByClass,
+		"find_text_by_tag":    h.findTextByTag,
 	})
 	L.Push(mod)
 	return 1
@@ -62,6 +65,36 @@ func (h *httpScrapeModule) findAttr(selector string, L *lua.LState) int {
 	return 2
 }
 
+func (h *httpScrapeModule) findAttrs(selector string, L *lua.LState) int {
+	body := L.ToString(1)
+	attrsCount := L.ToInt(2)
+	attrs := []string{}
+	for i := 1; i <= attrsCount; i++ {
+		attrNow := L.ToString(2 + i)
+		attrs = append(attrs, attrNow)
+	}
+	query := L.ToString(2 + attrsCount + 1)
+	L.Pop(2 + attrsCount + 1)
+	root, err := html.Parse(strings.NewReader(body))
+	if err != nil {
+		L.Push(lua.LNil)
+		L.Push(lua.LString(err.Error()))
+		return 2
+	}
+	results := scrape.FindAll(root, getMatcher(selector, query))
+	attrResults := []map[string]string{}
+	for _, result := range results {
+		attrResults = append(attrResults, make(map[string]string))
+		idx := len(attrResults) - 1
+		for _, attr := range attrs {
+			attrResults[idx][attr] = scrape.Attr(result, attr)
+		}
+	}
+	L.Push(luar.New(L, attrResults))
+	L.Push(lua.LNil)
+	return 2
+}
+
 func (h *httpScrapeModule) findText(selector string, L *lua.LState) int {
 	body := L.ToString(1)
 	query := L.ToString(2)
@@ -86,22 +119,34 @@ func (h *httpScrapeModule) findAttrByClass(L *lua.LState) int {
 	return h.findAttr("class", L)
 }
 
-func (h *httpScrapeModule) findTextByClass(L *lua.LState) int {
-	return h.findText("class", L)
-}
-
 func (h *httpScrapeModule) findAttrById(L *lua.LState) int {
 	return h.findAttr("id", L)
 }
 
-func (h *httpScrapeModule) findTextById(L *lua.LState) int {
-	return h.findText("id", L)
-}
-
 func (h *httpScrapeModule) findAttrByTag(L *lua.LState) int {
 	return h.findAttr("tag", L)
 }
 
+func (h *httpScrapeModule) findAttrsByClass(L *lua.LState) int {
+	return h.findAttrs("class", L)
+}
+
+func (h *httpScrapeModule) findAttrsById(L *lua.LState) int {
+	return h.findAttrs("id", L)
+}
+
+func (h *httpScrapeModule) findAttrsByTag(L *lua.LState) int {
+	return h.findAttrs("tag", L)
+}
+
+func (h *httpScrapeModule) findTextByClass(L *lua.LState) int {
+	return h.findText("class", L)
+}
+
+func (h *httpScrapeModule) findTextById(L *lua.LState) int {
+	return h.findText("id", L)
+}
+
 func (h *httpScrapeModule) findTextByTag(L *lua.LState) int {
 	return h.findText("tag", L)
 }

gluahttpscrape_test.go

@@ -20,45 +20,84 @@ func TestFindAttrByClass(t *testing.T) {
 	}
 }
 
-func TestFindTextByClass(t *testing.T) {
+func TestFindAttrById(t *testing.T) {
 	if err := evalLua(t, `
 		local scrape = require("scrape")
-		response, error = scrape.find_text_by_class("`+httpBody+`", "testclass")
-		assert_equal("My First Heading", response[1])
-		assert_equal("My First Heading", response[2])
+		response, error = scrape.find_attr_by_id("`+httpBody+`", "href", "testid")
+		assert_equal("testhref", response[1])
+		assert_equal("testhref2", response[2])
 	`); err != nil {
 		t.Errorf("Failed to evaluate script: %s", err)
 	}
 }
 
-func TestFindAttrById(t *testing.T) {
+func TestFindAttrByTag(t *testing.T) {
 	if err := evalLua(t, `
 		local scrape = require("scrape")
-		response, error = scrape.find_attr_by_id("`+httpBody+`", "href", "testid")
+		response, error = scrape.find_attr_by_tag("`+httpBody+`", "href", "h1")
 		assert_equal("testhref", response[1])
 		assert_equal("testhref2", response[2])
 	`); err != nil {
 		t.Errorf("Failed to evaluate script: %s", err)
 	}
 }
 
-func TestFindTextById(t *testing.T) {
+func TestFindAttrsByClass(t *testing.T) {
 	if err := evalLua(t, `
 		local scrape = require("scrape")
-		response, error = scrape.find_text_by_id("`+httpBody+`", "testid")
+		response, error = scrape.find_attrs_by_class("`+httpBody+`", 2, "id", "href", "testclass")
+		assert_equal("testhref", response[1]["href"])
+		assert_equal("testid", response[1]["id"])
+		assert_equal("testid", response[2]["id"])
+		assert_equal("testhref2", response[2]["href"])
+	`); err != nil {
+		t.Errorf("Failed to evaluate script: %s", err)
+	}
+}
+
+func TestFindAttrsById(t *testing.T) {
+	if err := evalLua(t, `
+		local scrape = require("scrape")
+		response, error = scrape.find_attrs_by_id("`+httpBody+`", 2, "id", "href", "testid")
+		assert_equal("testhref", response[1]["href"])
+		assert_equal("testid", response[1]["id"])
+		assert_equal("testid", response[2]["id"])
+		assert_equal("testhref2", response[2]["href"])
+	`); err != nil {
+		t.Errorf("Failed to evaluate script: %s", err)
+	}
+}
+
+func TestFindAttrsByTag(t *testing.T) {
+	if err := evalLua(t, `
+		local scrape = require("scrape")
+		response, error = scrape.find_attrs_by_tag("`+httpBody+`", 2, "id", "href", "h1")
+		assert_equal("testhref", response[1]["href"])
+		assert_equal("testid", response[1]["id"])
+		assert_equal("testid", response[2]["id"])
+		assert_equal("testhref2", response[2]["href"])
+	`); err != nil {
+		t.Errorf("Failed to evaluate script: %s", err)
+	}
+}
+
+func TestFindTextByClass(t *testing.T) {
+	if err := evalLua(t, `
+		local scrape = require("scrape")
+		response, error = scrape.find_text_by_class("`+httpBody+`", "testclass")
 		assert_equal("My First Heading", response[1])
 		assert_equal("My First Heading", response[2])
 	`); err != nil {
 		t.Errorf("Failed to evaluate script: %s", err)
 	}
 }
 
-func TestFindAttrByTag(t *testing.T) {
+func TestFindTextById(t *testing.T) {
 	if err := evalLua(t, `
 		local scrape = require("scrape")
-		response, error = scrape.find_attr_by_tag("`+httpBody+`", "href", "h1")
-		assert_equal("testhref", response[1])
-		assert_equal("testhref2", response[2])
+		response, error = scrape.find_text_by_id("`+httpBody+`", "testid")
+		assert_equal("My First Heading", response[1])
+		assert_equal("My First Heading", response[2])
 	`); err != nil {
 		t.Errorf("Failed to evaluate script: %s", err)
 	}