go-rod · ysmood · Feb 3, 2024 · Feb 1, 2024 · Feb 2, 2024 · Feb 2, 2024
diff --git a/examples_test.go b/examples_test.go
@@ -251,6 +251,21 @@ func Example_page_screenshot() {
 	_ = utils.OutputFile("my.jpg", img)
 }
 
+func Example_page_scroll_screenshot() {
+	browser := rod.New().MustConnect()
+
+	// capture entire browser viewport, returning jpg with quality=90
+	img, err := browser.MustPage("https://desktop.github.com/").MustWaitStable().ScrollScreenshot(&rod.ScrollScreenshotOptions{
+		Format:  proto.PageCaptureScreenshotFormatJpeg,
+		Quality: gson.Int(90),
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	_ = utils.OutputFile("my.jpg", img)
+}
+
 func Example_page_pdf() {
 	page := rod.New().MustConnect().MustPage("https://github.com").MustWaitLoad()
 

diff --git a/fixtures/scroll-y.html b/fixtures/scroll-y.html
@@ -0,0 +1,11 @@
+<!DOCTYPE html>
+<html>
+  <style>
+    button {
+      margin-top: 4500px;
+    }
+  </style>
+  <body>
+    <button>button</button>
+  </body>
+</html>
diff --git a/lib/utils/imageutil.go b/lib/utils/imageutil.go
@@ -0,0 +1,133 @@
+package utils
+
+import (
+	"bytes"
+	"fmt"
+	"image"
+	"image/jpeg"
+	"image/png"
+	"io"
+
+	"github.com/go-rod/rod/lib/proto"
+)
+
+// ImgWithBox is a image with a box, if the box is nil, it means the whole image.
+type ImgWithBox struct {
+	Img []byte
+	Box *image.Rectangle
+}
+
+// ImgOption is the option for image processing
+type ImgOption struct {
+	Quality int
+}
+
+// ImgProcessor is the interface for image processing
+type ImgProcessor interface {
+	Encode(img image.Image, opt *ImgOption) ([]byte, error)
+	Decode(file io.Reader) (image.Image, error)
+}
+
+type jpegProcessor struct{}
+
+func (p jpegProcessor) Encode(img image.Image, opt *ImgOption) ([]byte, error) {
+	var buf bytes.Buffer
+	var jpegOpt *jpeg.Options
+	if opt != nil {
+		jpegOpt = &jpeg.Options{Quality: opt.Quality}
+	}
+	err := jpeg.Encode(&buf, img, jpegOpt)
+	return buf.Bytes(), err
+}
+
+func (p jpegProcessor) Decode(file io.Reader) (image.Image, error) {
+	return jpeg.Decode(file)
+}
+
+type pngProcessor struct{}
+
+func (p pngProcessor) Encode(img image.Image, _ *ImgOption) ([]byte, error) {
+	var buf bytes.Buffer
+	err := png.Encode(&buf, img)
+	return buf.Bytes(), err
+}
+
+func (p pngProcessor) Decode(file io.Reader) (image.Image, error) {
+	return png.Decode(file)
+}
+
+// NewImgProcessor create a ImgProcessor by the format
+func NewImgProcessor(format proto.PageCaptureScreenshotFormat) (ImgProcessor, error) {
+	switch format {
+	case proto.PageCaptureScreenshotFormatJpeg:
+		return &jpegProcessor{}, nil
+	case "", proto.PageCaptureScreenshotFormatPng:
+		return &pngProcessor{}, nil
+	default:
+		return nil, fmt.Errorf("not support format: %v", format)
+	}
+}
+
+// SplicePngVertical splice png vertically, if there is only one image, it will return the image directly.
+// Only support png and jpeg format yet, webP is not supported because no suitable processing library was found in golang.
+func SplicePngVertical(files []ImgWithBox, format proto.PageCaptureScreenshotFormat, opt *ImgOption) ([]byte, error) {
+	if len(files) == 0 {
+		return nil, nil
+	}
+	if len(files) == 1 {
+		return files[0].Img, nil
+	}
+
+	var width, height int
+
+	processor, err := NewImgProcessor(format)
+	if err != nil {
+		return nil, err
+	}
+
+	var images []image.Image
+	for _, file := range files {
+		img, err := processor.Decode(bytes.NewReader(file.Img))
+		if err != nil {
+			return nil, err
+		}
+
+		images = append(images, img)
+		if file.Box != nil {
+			width = file.Box.Dx()
+			height += file.Box.Dy()
+		} else {
+			width = img.Bounds().Dx()
+			height += img.Bounds().Dy()
+		}
+	}
+
+	spliceImg := image.NewRGBA(image.Rect(0, 0, width, height))
+
+	var destY int
+	for i, file := range files {
+		img := images[i]
+		bounds := img.Bounds()
+
+		if file.Box != nil {
+			bounds = *file.Box
+		}
+		start := bounds.Min
+		end := bounds.Max
+		for y := start.Y; y < end.Y; y++ {
+			for x := start.X; x < end.X; x++ {
+				color := img.At(x, y)
+				spliceImg.Set(x, y-start.Y+destY, color)
+			}
+		}
+
+		destY += bounds.Dy()
+	}
+
+	bs, err := processor.Encode(spliceImg, opt)
+	if err != nil {
+		return nil, err
+	}
+
+	return bs, nil
+}
diff --git a/lib/utils/imageutil_test.go b/lib/utils/imageutil_test.go
@@ -0,0 +1,207 @@
+package utils
+
+import (
+	"bytes"
+	"image"
+	"testing"
+
+	"github.com/go-rod/rod/lib/proto"
+	"github.com/ysmood/got"
+)
+
+var setup = got.Setup(nil)
+
+func TestSplicePngVertical(t *testing.T) {
+	g := setup(t)
+	a := image.NewRGBA(image.Rect(0, 0, 1000, 200))
+	b := image.NewRGBA(image.Rect(0, 0, 1000, 300))
+
+	t.Run("jpeg", func(t *testing.T) {
+		format := proto.PageCaptureScreenshotFormatJpeg
+		processor, err := NewImgProcessor(format)
+		if err != nil {
+			g.Err(err)
+		}
+		aBs, _ := processor.Encode(a, nil)
+		bBs, _ := processor.Encode(b, nil)
+
+		bs, err := SplicePngVertical([]ImgWithBox{
+			{Img: aBs},
+			{Img: bBs},
+		}, format, nil)
+		g.E(err)
+
+		img, err := processor.Decode(bytes.NewBuffer(bs))
+		g.E(err)
+
+		g.Eq(img.Bounds().Dy(), 500)
+		g.Eq(img.Bounds().Dx(), 1000)
+	})
+	t.Run("jpegWithOptions", func(t *testing.T) {
+		format := proto.PageCaptureScreenshotFormatJpeg
+		processor, err := NewImgProcessor(format)
+		g.E(err)
+
+		aBs, _ := processor.Encode(a, nil)
+		bBs, _ := processor.Encode(b, nil)
+
+		bs, err := SplicePngVertical([]ImgWithBox{
+			{Img: aBs},
+			{Img: bBs},
+		}, format, &ImgOption{
+			Quality: 10,
+		})
+		g.E(err)
+
+		img, err := processor.Decode(bytes.NewBuffer(bs))
+		g.E(err)
+
+		g.Eq(img.Bounds().Dy(), 500)
+		g.Eq(img.Bounds().Dx(), 1000)
+	})
+	t.Run("jpegWithBox", func(t *testing.T) {
+		format := proto.PageCaptureScreenshotFormatJpeg
+		processor, err := NewImgProcessor(format)
+		g.E(err)
+
+		aBs, _ := processor.Encode(a, nil)
+		bBs, _ := processor.Encode(b, nil)
+
+		bs, err := SplicePngVertical([]ImgWithBox{
+			{
+				Img: aBs,
+				Box: &image.Rectangle{
+					Max: image.Point{
+						X: a.Bounds().Dx(),
+						Y: 100,
+					},
+				},
+			},
+			{Img: bBs},
+		}, format, nil)
+		g.E(err)
+
+		img, err := processor.Decode(bytes.NewBuffer(bs))
+		g.E(err)
+
+		g.Eq(img.Bounds().Dy(), 400)
+		g.Eq(img.Bounds().Dx(), 1000)
+	})
+	t.Run("errorEncode", func(t *testing.T) {
+		format := proto.PageCaptureScreenshotFormatPng
+		processor, err := NewImgProcessor(format)
+		g.E(err)
+
+		aBs, _ := processor.Encode(a, nil)
+		bBs, _ := processor.Encode(b, nil)
+
+		_, err = SplicePngVertical([]ImgWithBox{
+			{
+				Img: aBs,
+				Box: &image.Rectangle{},
+			},
+			{
+				Img: bBs,
+				Box: &image.Rectangle{},
+			},
+		}, format, nil)
+		// invalid image size: 0x0
+		g.Err(err)
+	})
+	t.Run("noFile", func(t *testing.T) {
+		_, err := SplicePngVertical(nil, "", nil)
+		g.E(err)
+	})
+	t.Run("oneFile", func(t *testing.T) {
+		bs, err := SplicePngVertical([]ImgWithBox{
+			{Img: []byte{1}},
+		}, "", nil)
+		g.E(err)
+		g.Eq(1, len(bs))
+	})
+	t.Run("unsupportedFormat", func(t *testing.T) {
+		_, err := SplicePngVertical([]ImgWithBox{
+			{Img: []byte{1}},
+			{Img: []byte{1}},
+		}, "gif", nil)
+		g.Err(err)
+	})
+	t.Run("errorFile", func(t *testing.T) {
+		_, err := SplicePngVertical([]ImgWithBox{
+			{Img: []byte{1}},
+			{Img: []byte{1}},
+		}, "", nil)
+		g.Err(err)
+	})
+}
+
+func TestNewImgProcessor(t *testing.T) {
+	g := setup(t)
+	type args struct {
+		format proto.PageCaptureScreenshotFormat
+	}
+	tests := []struct {
+		name    string
+		args    args
+		wantErr bool
+	}{
+		{
+			name: "jpeg",
+			args: args{
+				format: proto.PageCaptureScreenshotFormatJpeg,
+			},
+			wantErr: false,
+		},
+		{
+			name: "default",
+			args: args{
+				format: "",
+			},
+			wantErr: false,
+		},
+		{
+			name: "png",
+			args: args{
+				format: proto.PageCaptureScreenshotFormatPng,
+			},
+			wantErr: false,
+		},
+		{
+			name: "webP",
+			args: args{
+				/* cspell: disable-next-line */
+				format: proto.PageCaptureScreenshotFormatWebp,
+			},
+			wantErr: true,
+		},
+	}
+
+	a := image.NewRGBA(image.Rect(0, 0, 1000, 200))
+	// errImg := image.NewRGBA(image.Rect(0, 0, 0, 0))
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			processor, err := NewImgProcessor(tt.args.format)
+			if tt.wantErr {
+				g.Eq(err != nil, tt.wantErr)
+			}
+			if err != nil {
+				return
+			}
+			buf, err := processor.Encode(a, nil)
+			if err != nil {
+				g.Err(err)
+			}
+			img, err := processor.Decode(bytes.NewBuffer(buf))
+			if err != nil {
+				g.Err(err)
+			}
+
+			g.Eq(1000, img.Bounds().Dx())
+			g.Eq(200, img.Bounds().Dy())
+
+			_, err = processor.Decode(bytes.NewBuffer(nil))
+			g.Err(err)
+		})
+	}
+}
diff --git a/must.go b/must.go
@@ -388,6 +388,15 @@ func (p *Page) MustScreenshotFullPage(toFile ...string) []byte {
 	return bin
 }
 
+// MustScrollScreenshotPage is similar to [Page.ScrollScreenshot].
+// If the toFile is "", it Page.will save output to "tmp/screenshots" folder, time as the file name.
+func (p *Page) MustScrollScreenshotPage(toFile ...string) []byte {
+	bin, err := p.ScrollScreenshot(nil)
+	p.e(err)
+	p.e(saveFile(saveFileTypeScreenshot, bin, toFile))
+	return bin
+}
+
 // MustPDF is similar to [Page.PDF].
 // If the toFile is "", it Page.will save output to "tmp/pdf" folder, time as the file name.
 func (p *Page) MustPDF(toFile ...string) []byte {