tren/extract/translation.go

package extract

import (
	"fmt"
	"strings"

	"github.com/PuerkitoBio/goquery"
)

type Entry struct {
  Text string
  Detail string
}

type Translation struct {
  Suggestions []string
  Fields []string
  Entries [][]Entry
}

func (t *Translation) getSuggestions(doc *goquery.Document) error {
  doc.Find(".suggestion-list li").Each(func(_ int, sel *goquery.Selection){
    if link := sel.Find("a"); link != nil {
      t.Suggestions = append(t.Suggestions, link.Text())
    }
  })

  return nil
}

func (t *Translation) getEntry(tds *goquery.Selection) {
  if tds == nil || tds.Length() < 4{
    return
  }

  var (
    entries []Entry
    index int = 0
  )

  tds.Each(func(_ int,sel *goquery.Selection){
    if index < 1 || index > 3 {
      index++
      return
    }

    link := sel.Find("a")
    detail := sel.Find("i")
    text := ""

    if link.Length() == 0 {
      text = sel.Text()
    } else {
      text = link.Text()
    }

    if detail == nil {
      entries = append(entries, Entry{
        Text: text,
      })
    }else {
      entries = append(entries, Entry{
        Text: text,
        Detail: detail.Text(),
      })
    }

    index++
  })

  t.Entries = append(t.Entries, entries)
}

func (t *Translation) getEntries(trs *goquery.Selection) {
  trs.Each(func(_ int, sel *goquery.Selection){
    tds := sel.Find("td")

    if tds == nil {
      return
    }

    if class, ok := tds.Last().Attr("class"); !ok {
      return
    }else if !strings.Contains(class, "rc4 hidden-xs"){
      return
    }

    t.getEntry(tds)
  })
}

func (t *Translation) GetResults(doc *goquery.Document) error {
  table := doc.Find(".searchResultsTable")

  if table.Length() == 0 {
    return t.getSuggestions(doc)
  }

  trs := table.Find("tbody tr")

  if trs.Length() == 0 {
    return fmt.Errorf("failed to get table body")
  }

  trs.First().Find("th").Each(func(_ int, sel *goquery.Selection){
    if sel.Text() != "" {
      t.Fields = append(t.Fields, sel.Text())
    }
  })

  t.getEntries(trs)
  return nil
}