add WithCancel shortcut

browser_test.go

@@ -130,9 +130,9 @@ func (s *S) TestMonitor() {
 	defer b.MustClose()
 	p := b.MustPage(srcFile("fixtures/click.html")).MustWaitLoad()
 
-	ctx, cancel := context.WithCancel(context.Background())
+	b, cancel := b.WithCancel()
 	defer cancel()
-	host := b.Context(ctx).ServeMonitor("")
+	host := b.ServeMonitor("")
 
 	page := s.page.MustNavigate(host)
 	s.Contains(page.MustElement("#targets a").MustParent().MustHTML(), string(p.TargetID))

context.go

@@ -33,6 +33,12 @@ func (b *Browser) CancelTimeout() *Browser {
 	return b.Context(val.parent)
 }
 
+// WithCancel returns a clone with a context cancel function
+func (b *Browser) WithCancel() (*Browser, func()) {
+	ctx, cancel := context.WithCancel(b.ctx)
+	return b.Context(ctx), cancel
+}
+
 // Sleeper for chained sub-operations
 func (b *Browser) Sleeper(sleeper func() utils.Sleeper) *Browser {
 	newObj := *b
@@ -60,6 +66,12 @@ func (p *Page) CancelTimeout() *Page {
 	return p.Context(val.parent)
 }
 
+// WithCancel returns a clone with a context cancel function
+func (p *Page) WithCancel() (*Page, func()) {
+	ctx, cancel := context.WithCancel(p.ctx)
+	return p.Context(ctx), cancel
+}
+
 // Sleeper for chained sub-operations
 func (p *Page) Sleeper(sleeper func() utils.Sleeper) *Page {
 	newObj := *p
@@ -87,6 +99,12 @@ func (el *Element) CancelTimeout() *Element {
 	return el.Context(val.parent)
 }
 
+// WithCancel returns a clone with a context cancel function
+func (el *Element) WithCancel() (*Element, func()) {
+	ctx, cancel := context.WithCancel(el.ctx)
+	return el.Context(ctx), cancel
+}
+
 // Sleeper for chained sub-operations
 func (el *Element) Sleeper(sleeper func() utils.Sleeper) *Element {
 	newObj := *el

element_test.go

@@ -139,6 +139,8 @@ func (s *S) TestMouseMoveErr() {
 func (s *S) TestElementContext() {
 	p := s.page.MustNavigate(srcFile("fixtures/click.html"))
 	el := p.MustElement("button").Timeout(time.Hour).CancelTimeout()
+	el, cancel := el.WithCancel()
+	defer cancel()
 	el.Sleeper(rod.DefaultSleeper).MustClick()
 }
 

examples_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math/rand"
 	"time"
 
 	"github.com/go-rod/rod"
@@ -129,9 +130,12 @@ func Example_timeout_handling() {
 	}
 	{
 		// Use this way you can customize your own way to cancel long-running task
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		page.Context(ctx).MustElement("a")
-		cancel()
+		page, cancel := page.WithCancel()
+		go func() {
+			time.Sleep(time.Duration(rand.Int())) // cancel after randomly time
+			cancel()
+		}()
+		page.MustElement("a")
 	}
 }
 

hijack.go

@@ -375,9 +375,9 @@ func (p *Page) GetDownloadFile(pattern string, resourceType proto.NetworkResourc
 
 	r := p.HijackRequests()
 
-	ctx, cancel := context.WithCancel(p.ctx)
+	p, cancel := p.WithCancel()
 	downloading := &proto.PageDownloadWillBegin{}
-	waitDownload := p.Context(ctx).WaitEvent(downloading)
+	waitDownload := p.WaitEvent(downloading)
 
 	return func() (http.Header, []byte, error) {
 		defer enable()

page.go

@@ -235,10 +235,10 @@ func (p *Page) Close() error {
 	}
 
 	success := true
-	ctx, cancel := context.WithCancel(p.ctx)
+	p, cancel := p.WithCancel()
 	defer cancel()
 
-	wait := p.Context(ctx).EachEvent(func(e *proto.TargetDetachedFromTarget) bool {
+	wait := p.EachEvent(func(e *proto.TargetDetachedFromTarget) bool {
 		return e.TargetID == e.TargetID
 	}, func(e *proto.PageJavascriptDialogClosed) bool {
 		success = e.Result
@@ -328,17 +328,15 @@ func (p *Page) PDF(req *proto.PagePrintToPDF) (*StreamReader, error) {
 
 // WaitOpen waits for the next new page opened by the current one
 func (p *Page) WaitOpen() func() (*Page, error) {
-	b := p.browser.Context(p.ctx)
 	var targetID proto.TargetTargetID
 
-	ctx, cancel := context.WithCancel(p.ctx)
-	wait := b.Context(ctx).EachEvent(func(e *proto.TargetTargetCreated) bool {
+	b := p.browser.Context(p.ctx)
+	wait := b.EachEvent(func(e *proto.TargetTargetCreated) bool {
 		targetID = e.TargetInfo.TargetID
 		return e.TargetInfo.OpenerID == p.TargetID
 	})
 
 	return func() (*Page, error) {
-		defer cancel()
 		wait()
 		return b.PageFromTarget(targetID)
 	}
@@ -502,7 +500,7 @@ func (p *Page) Expose(name string) (callback chan string, stop func(), err error
 	}
 
 	callback = make(chan string)
-	ctx, cancel := context.WithCancel(p.ctx)
+	p, cancel := p.WithCancel()
 	stop = func() {
 		cancel()
 		_ = proto.RuntimeRemoveBinding{Name: name}.Call(p)
@@ -511,7 +509,7 @@ func (p *Page) Expose(name string) (callback chan string, stop func(), err error
 	go p.EachEvent(func(e *proto.RuntimeBindingCalled) bool {
 		if e.Name == name {
 			select {
-			case <-ctx.Done():
+			case <-p.ctx.Done():
 				return true
 			case callback <- e.Payload:
 			}