add Element.Screenshot

twolao · Mar 9, 2020 · 09ee7c4 · 09ee7c4
1 parent 1cf4a66
commit 09ee7c4
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 3 deletions.
diff --git a/element.go b/element.go
@@ -338,6 +338,41 @@ func (el *Element) ResourceE() ([]byte, error) {
 	return bin, nil
 }
 
+// ScreenshotE of the area of the element
+func (el *Element) ScreenshotE(format string, quality int) ([]byte, error) {
+	err := el.WaitVisibleE()
+	if err != nil {
+		return nil, err
+	}
+
+	err = el.ScrollIntoViewIfNeededE()
+	if err != nil {
+		return nil, err
+	}
+
+	box, err := el.BoxE()
+	if err != nil {
+		return nil, err
+	}
+
+	opts := cdp.Object{
+		"format": format,
+		"clip": cdp.Object{
+			"x":      box.Left,
+			"y":      box.Top,
+			"width":  box.Width,
+			"height": box.Height,
+			"scale":  1,
+		},
+	}
+
+	if quality > -1 {
+		opts["quality"] = quality
+	}
+
+	return el.page.Root().ScreenshotE(opts)
+}
+
 // ReleaseE doc is the same as the method Release
 func (el *Element) ReleaseE() error {
 	return el.page.Context(el.ctx).ReleaseE(el.ObjectID)

diff --git a/element_test.go b/element_test.go
@@ -1,8 +1,10 @@
 package rod_test
 
 import (
+	"bytes"
 	"context"
 	"errors"
+	"image/png"
 	"path/filepath"
 	"time"
 
@@ -164,6 +166,16 @@ func (s *S) TestResource() {
 	s.Equal(15148, len(p.Element("img").Resource()))
 }
 
+func (s *S) TestElementScreenshot() {
+	p := s.page.Navigate(srcFile("fixtures/click.html"))
+
+	data := p.Element("h4").Screenshot()
+	img, err := png.Decode(bytes.NewBuffer(data))
+	kit.E(err)
+	s.EqualValues(200, img.Bounds().Dx())
+	s.EqualValues(30, img.Bounds().Dy())
+}
+
 func (s *S) TestUseReleasedElement() {
 	p := s.page.Navigate(srcFile("fixtures/click.html"))
 	btn := p.Element("button")

diff --git a/fixtures/click.html b/fixtures/click.html
@@ -1,5 +1,9 @@
 <html>
     <style>
+        h4 {
+            width: 200px;
+            height: 30px;
+        }
         button {
             margin: 100px;
         }

diff --git a/page_test.go b/page_test.go
@@ -216,13 +216,12 @@ func (s *S) TestPagePause() {
 
 func (s *S) TestPageScreenshot() {
 	p := s.page.Navigate(srcFile("fixtures/click.html"))
-	p.Viewport(400, 300, 1, false)
 	p.Element("button")
 	data := p.Screenshot()
 	img, err := png.Decode(bytes.NewBuffer(data))
 	kit.E(err)
-	s.Equal(400, img.Bounds().Dx())
-	s.Equal(300, img.Bounds().Dy())
+	s.Equal(800, img.Bounds().Dx())
+	s.Equal(600, img.Bounds().Dy())
 }
 
 func (s *S) TestPageTraceDir() {

diff --git a/rod_test.go b/rod_test.go
@@ -61,6 +61,7 @@ func Test(t *testing.T) {
 	defer s.browser.Close()
 
 	s.page = s.browser.Page("")
+	s.page.Viewport(800, 600, 1, false)
 
 	suite.Run(t, s)
 }
diff --git a/sugar.go b/sugar.go
@@ -464,6 +464,13 @@ func (el *Element) Resource() []byte {
 	return bin
 }
 
+// Screenshot of the area of the element
+func (el *Element) Screenshot() []byte {
+	bin, err := el.ScreenshotE("png", -1)
+	kit.E(err)
+	return bin
+}
+
 // Release remote object on browser
 func (el *Element) Release() {
 	kit.E(el.ReleaseE())