package builtins import ( "io" "strings" "github.com/PuerkitoBio/goquery" "go.starlark.net/starlark" "go.starlark.net/starlarkstruct" ) var ( _ starlark.Iterable = (*starlarkSelection)(nil) _ starlark.Sliceable = (*starlarkSelection)(nil) _ starlark.Sequence = (*starlarkSelection)(nil) _ starlark.Value = (*starlarkSelection)(nil) ) var htmlModule = &starlarkstruct.Module{ Name: "html", Members: starlark.StringDict{ "parse": starlark.NewBuiltin("html.parse", htmlParse), }, } func htmlParse(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var val starlark.Value err := starlark.UnpackArgs("html.selection.find", args, kwargs, "from", &val) if err != nil { return nil, err } var r io.ReadCloser switch val := val.(type) { case starlark.String: r = io.NopCloser(strings.NewReader(string(val))) case starlark.Bytes: r = io.NopCloser(strings.NewReader(string(val))) case starlarkReader: r = val } defer r.Close() doc, err := goquery.NewDocumentFromReader(r) if err != nil { return nil, err } return newStarlarkSelection(doc.Selection), nil } type starlarkSelection struct { sel *goquery.Selection *starlarkstruct.Struct } func newStarlarkSelection(sel *goquery.Selection) starlarkSelection { ss := starlarkSelection{sel: sel} ss.Struct = starlarkstruct.FromStringDict(starlark.String("html.selection"), starlark.StringDict{ "text": starlark.NewBuiltin("html.selection.text", ss.text), "html": starlark.NewBuiltin("html.selection.html", ss.html), "children": starlark.NewBuiltin("html.selection.children", ss.children), "parent": starlark.NewBuiltin("html.selection.parent", ss.parent), "find": starlark.NewBuiltin("html.selection.find", ss.find), "add": starlark.NewBuiltin("html.selection.add", ss.add), "attr": starlark.NewBuiltin("html.selection.attr", ss.attr), "has_class": starlark.NewBuiltin("html.selection.has_class", ss.hasClass), "index_selector": starlark.NewBuiltin("html.selection.index_selector", ss.indexSelector), "and_self": starlark.NewBuiltin("html.selection.and_self", ss.andSelf), "first": starlark.NewBuiltin("html.selection.first", ss.first), "last": starlark.NewBuiltin("html.selection.last", ss.last), "next": starlark.NewBuiltin("html.selection.last", ss.next), "next_all": starlark.NewBuiltin("html.selection.next_all", ss.nextAll), "next_until": starlark.NewBuiltin("html.selection.next_until", ss.nextUntil), "prev": starlark.NewBuiltin("html.selection.prev", ss.prev), "prev_all": starlark.NewBuiltin("html.selection.prev_all", ss.prevAll), "prev_until": starlark.NewBuiltin("html.selection.prev_until", ss.prevUntil), }) return ss } func (ss starlarkSelection) Truth() starlark.Bool { return len(ss.sel.Nodes) > 0 } func (ss starlarkSelection) Len() int { return ss.sel.Length() } func (ss starlarkSelection) Index(i int) starlark.Value { return newStarlarkSelection(ss.sel.Slice(i, i+1)) } func (ss starlarkSelection) Slice(start, end, _ int) starlark.Value { return newStarlarkSelection(ss.sel.Slice(start, end)) } func (ss starlarkSelection) Iterate() starlark.Iterator { return newSelectionIterator(ss.sel) } func (ss starlarkSelection) text(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { return starlark.String(ss.sel.Text()), nil } func (ss starlarkSelection) html(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { s, err := ss.sel.Html() return starlark.String(s), err } func (ss starlarkSelection) children(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { return newStarlarkSelection(ss.sel.Children()), nil } func (ss starlarkSelection) parent(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { return newStarlarkSelection(ss.sel.Parent()), nil } func (ss starlarkSelection) find(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var selector string err := starlark.UnpackArgs("html.selection.find", args, kwargs, "selector", &selector) if err != nil { return nil, err } return newStarlarkSelection(ss.sel.Find(selector)), nil } func (ss starlarkSelection) add(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var selector string err := starlark.UnpackArgs("html.selection.add", args, kwargs, "selector", &selector) if err != nil { return nil, err } return newStarlarkSelection(ss.sel.Add(selector)), nil } func (ss starlarkSelection) indexSelector(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var selector string err := starlark.UnpackArgs("html.selection.index_selector", args, kwargs, "selector", &selector) if err != nil { return nil, err } return starlark.MakeInt(ss.sel.IndexSelector(selector)), nil } func (ss starlarkSelection) attr(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var name, def string err := starlark.UnpackArgs("html.selection.find", args, kwargs, "name", &name, "default??", &def) if err != nil { return nil, err } return starlark.String(ss.sel.AttrOr(name, def)), nil } func (ss starlarkSelection) hasClass(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var name string err := starlark.UnpackArgs("html.selection.has_class", args, kwargs, "name", &name) if err != nil { return nil, err } return starlark.Bool(ss.sel.HasClass(name)), nil } func (ss starlarkSelection) andSelf(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { return newStarlarkSelection(ss.sel.AndSelf()), nil } func (ss starlarkSelection) first(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { return newStarlarkSelection(ss.sel.First()), nil } func (ss starlarkSelection) last(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { return newStarlarkSelection(ss.sel.Last()), nil } func (ss starlarkSelection) next(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { return newStarlarkSelection(ss.sel.Next()), nil } func (ss starlarkSelection) nextAll(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { return newStarlarkSelection(ss.sel.NextAll()), nil } func (ss starlarkSelection) nextUntil(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var selector string err := starlark.UnpackArgs("html.selection.next_until", args, kwargs, "selector", &selector) if err != nil { return nil, err } return newStarlarkSelection(ss.sel.NextUntil(selector)), nil } func (ss starlarkSelection) prev(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { return newStarlarkSelection(ss.sel.Prev()), nil } func (ss starlarkSelection) prevAll(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { return newStarlarkSelection(ss.sel.PrevAll()), nil } func (ss starlarkSelection) prevUntil(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var selector string err := starlark.UnpackArgs("html.selection.prev_until", args, kwargs, "selector", &selector) if err != nil { return nil, err } return newStarlarkSelection(ss.sel.PrevUntil(selector)), nil } type starlarkSelectionIterator struct { sel *goquery.Selection index int } func newSelectionIterator(sel *goquery.Selection) *starlarkSelectionIterator { return &starlarkSelectionIterator{sel: sel} } func (ssi *starlarkSelectionIterator) Next(v *starlark.Value) bool { if ssi.index == ssi.sel.Length() { return false } *v = newStarlarkSelection(ssi.sel.Slice(ssi.index, ssi.index+1)) ssi.index++ return true } func (ssi *starlarkSelectionIterator) Done() {}