Some useful helpers for chromedp.

cdputil.go

// Package cdputil contains some useful helpers for chromedp.
package cdputil

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/chromedp/cdproto/cdp"
	"github.com/chromedp/cdproto/runtime"
	"github.com/chromedp/chromedp"
)

// Group groups multiple actions into a single action.
func Group(actions ...chromedp.Action) chromedp.Action {
	return chromedp.Tasks(actions)
}

// Timeout specifies a timeout for the provided action.
func Timeout(timeout time.Duration, action chromedp.Action) chromedp.Action {
	return chromedp.ActionFunc(func(ctx context.Context) error {
		ctx, cancel := context.WithCancel(ctx)
		defer cancel()

		var us bool
		t := time.AfterFunc(timeout, func() {
			us = true
			cancel()
		})

		err := action.Do(ctx)
		t.Stop()

		if us && errors.Is(err, context.Canceled) {
			err = xerr{
				is:      context.DeadlineExceeded,
				message: fmt.Sprintf("timed out after %s", timeout),
			}
		}
		return err
	})
}

// ElementCount executes the function with the number of matching nodes for the
// provided selector and runs the returned action.
func ElementCount(sel interface{}, fn func(ctx context.Context, n int) chromedp.Action) chromedp.Action {
	return chromedp.QueryAfter(sel, func(ctx context.Context, eci runtime.ExecutionContextID, nodes ...*cdp.Node) error {
		if a := fn(ctx, len(nodes)); a != nil {
			return a.Do(ctx)
		}
		return nil
	}, chromedp.AtLeast(0))
}

// ElementTernary executes action if the provided selector matches, or
// actionMissing otherwise.
func ElementTernary(sel interface{}, action, actionMissing chromedp.Action) chromedp.Action {
	return ElementCount(sel, func(ctx context.Context, n int) chromedp.Action {
		if n == 0 {
			return actionMissing
		} else {
			return action
		}
	})
}

// Error returns an error.
func ErrorAlways(err error) chromedp.Action {
	return chromedp.ActionFunc(func(ctx context.Context) error {
		return err
	})
}

// Errorf returns an error.
func ErrorAlwaysf(is error, format string, a ...interface{}) chromedp.Action {
	return chromedp.ActionFunc(func(ctx context.Context) error {
		return xerr{
			is:          is,
			wrapped:     fmt.Errorf(format, a...),
			showWrapped: true,
		}
	})
}

// ErrorWrap wraps error messages returned by the action with the
// provided message, while also making errors.Is match on is if provided.
func ErrorWrap(msg string, is error, action chromedp.Action) chromedp.Action {
	return chromedp.ActionFunc(func(ctx context.Context) error {
		err := action.Do(ctx)
		if err != nil {
			err = xerr{
				is:          is,
				wrapped:     err,
				message:     msg,
				showWrapped: true,
			}
		}
		return err
	})
}

// ErrorNodeText returns an error in the format `msg: "<text of sel>"`,
// matching errors.Is for is if provided. If sel is not found, a generic error
// message is used instead.
func ErrorNodeText(msg string, is error, sel interface{}) chromedp.Action {
	return chromedp.ActionFunc(func(ctx context.Context) error {
		var txt string
		if err := chromedp.Text(sel, &txt, chromedp.AtLeast(0)).Do(ctx); err != nil {
			return xerr{
				is:          is,
				wrapped:     fmt.Errorf("unknown error (failed to extract message from %q: %w)", sel, err),
				message:     msg,
				showWrapped: true,
			}
		}
		txt = strings.TrimSpace(txt)
		if len(txt) == 0 {
			return xerr{
				is:          is,
				wrapped:     fmt.Errorf("unknown error (failed to extract message from %q: no match)", sel),
				message:     msg,
				showWrapped: true,
			}
		}
		return xerr{
			is:          is,
			wrapped:     fmt.Errorf("%q", txt),
			message:     msg,
			showWrapped: true,
		}
	})
}

// ErrorNoSelectorMatch returns an error if any of the provided selectors do not
// match one or more nodes. Note that you can use commas to OR selectors.
func ErrorNoSelectorMatch(sels ...interface{}) chromedp.QueryAction {
	tasks := make(chromedp.Tasks, len(sels))
	for i, sel := range sels {
		tasks[i] = ElementTernary(sel, nil, ErrorAlways(fmt.Errorf("no match for selector %#v", sel)))
	}
	return tasks
}

type screenshotterKey struct{}
type ScreenshotterHandlerFunc func(name string, png []byte)

// ScreenshotterHandler appends a handler for screenshots.
func ScreenshotterHandler(ctx context.Context, handler ScreenshotterHandlerFunc) context.Context {
	if fns, ok := ctx.Value(screenshotterKey{}).([]ScreenshotterHandlerFunc); ok {
		return context.WithValue(ctx, screenshotterKey{}, append([]ScreenshotterHandlerFunc{handler}, fns...))
	}
	return context.WithValue(ctx, screenshotterKey{}, []ScreenshotterHandlerFunc{handler})
}

// ScreenshotterFull adds a chromedp.FullScreenshot with the provided name to
// all ScreenshotterHandlerFuncs registered on the context.
func ScreenshotterFull(name string) chromedp.Action {
	return chromedp.ActionFunc(func(ctx context.Context) error {
		if fns, ok := ctx.Value(screenshotterKey{}).([]ScreenshotterHandlerFunc); ok {
			var png []byte
			if err := chromedp.FullScreenshot(&png, 100).Do(ctx); err != nil {
				return err
			}
			for _, fn := range fns {
				fn(name, png)
			}
		}
		return nil
	})
}

// xerr is a versatile error type which allows for complex options like having
// an error wrap another with a custom message while matching errors.Is for a
// different error, making another error match errors.Is for an unrelated one
// out of the chain, adding a message to another error, allowing an error to be
// unwrapped while hiding the wrapped message, and so on.
type xerr struct {
	is          error
	wrapped     error
	message     string
	showWrapped bool
}

func (err xerr) Is(target error) bool {
	if err.is != nil && errors.Is(err.is, target) {
		return true
	}
	if err.wrapped != nil && errors.Is(err.wrapped, target) {
		return true
	}
	return false
}

func (err xerr) Unwrap() error {
	if err.wrapped == nil {
		return err.is
	}
	return err.wrapped
}

func (err xerr) Error() string {
	if err.wrapped == nil || !err.showWrapped {
		if err.message != "" {
			return err.message
		} else if err.is != nil {
			return err.is.Error()
		} else {
			return "unknown error"
		}
	}
	if err.message == "" {
		if err.is != nil {
			return fmt.Sprintf("%v: %v", err.is, err.wrapped)
		} else {
			return err.wrapped.Error()
		}
	}
	return fmt.Sprintf("%s: %v", err.message, err.wrapped)
}