From 42f356b4d86ea0b3f3f2388e567d999aee72dab9 Mon Sep 17 00:00:00 2001
From: Kenneth Shaw <kenshaw@gmail.com>
Date: Fri, 2 Feb 2024 03:08:30 +0700
Subject: [PATCH] Adding ScreenshotNodes

---
 go.mod        |  4 ++--
 go.sum        |  8 ++++----
 screenshot.go | 51 +++++++++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/go.mod b/go.mod
index 957523e6..5be0af14 100644
--- a/go.mod
+++ b/go.mod
@@ -3,8 +3,8 @@ module github.com/chromedp/chromedp
 go 1.16
 
 require (
-	github.com/chromedp/cdproto v0.0.0-20231011050154-1d073bb38998
-	github.com/gobwas/ws v1.3.0
+	github.com/chromedp/cdproto v0.0.0-20240127002248-bd7a66284627
+	github.com/gobwas/ws v1.3.2
 	github.com/ledongthuc/pdf v0.0.0-20220302134840-0c2507a12d80
 	github.com/mailru/easyjson v0.7.7
 	github.com/orisano/pixelmatch v0.0.0-20220722002657-fb0b55479cde
diff --git a/go.sum b/go.sum
index 868a2ba8..007164c7 100644
--- a/go.sum
+++ b/go.sum
@@ -1,13 +1,13 @@
-github.com/chromedp/cdproto v0.0.0-20231011050154-1d073bb38998 h1:2zipcnjfFdqAjOQa8otCCh0Lk1M7RBzciy3s80YAKHk=
-github.com/chromedp/cdproto v0.0.0-20231011050154-1d073bb38998/go.mod h1:GKljq0VrfU4D5yc+2qA6OVr8pmO/MBbPEWqWQ/oqGEs=
+github.com/chromedp/cdproto v0.0.0-20240127002248-bd7a66284627 h1:L5rJ/yzLfSU3kcjsjq11xYDqAdianisL21CXQ/08Zag=
+github.com/chromedp/cdproto v0.0.0-20240127002248-bd7a66284627/go.mod h1:GKljq0VrfU4D5yc+2qA6OVr8pmO/MBbPEWqWQ/oqGEs=
 github.com/chromedp/sysutil v1.0.0 h1:+ZxhTpfpZlmchB58ih/LBHX52ky7w2VhQVKQMucy3Ic=
 github.com/chromedp/sysutil v1.0.0/go.mod h1:kgWmDdq8fTzXYcKIBqIYvRRTnYb9aNS9moAV0xufSww=
 github.com/gobwas/httphead v0.1.0 h1:exrUm0f4YX0L7EBwZHuCF4GDp8aJfVeBrlLQrs6NqWU=
 github.com/gobwas/httphead v0.1.0/go.mod h1:O/RXo79gxV8G+RqlR/otEwx4Q36zl9rqC5u12GKvMCM=
 github.com/gobwas/pool v0.2.1 h1:xfeeEhW7pwmX8nuLVlqbzVc7udMDrwetjEv+TZIz1og=
 github.com/gobwas/pool v0.2.1/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=
-github.com/gobwas/ws v1.3.0 h1:sbeU3Y4Qzlb+MOzIe6mQGf7QR4Hkv6ZD0qhGkBFL2O0=
-github.com/gobwas/ws v1.3.0/go.mod h1:hRKAFb8wOxFROYNsT1bqfWnhX+b5MFeJM9r2ZSwg/KY=
+github.com/gobwas/ws v1.3.2 h1:zlnbNHxumkRvfPWgfXu8RBwyNR1x8wh9cf5PTOCqs9Q=
+github.com/gobwas/ws v1.3.2/go.mod h1:hRKAFb8wOxFROYNsT1bqfWnhX+b5MFeJM9r2ZSwg/KY=
 github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY=
 github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=
 github.com/ledongthuc/pdf v0.0.0-20220302134840-0c2507a12d80 h1:6Yzfa6GP0rIo/kULo2bwGEkFvCePZ3qHDDTC3/J9Swo=
diff --git a/screenshot.go b/screenshot.go
index 3ed0c798..0130edbe 100644
--- a/screenshot.go
+++ b/screenshot.go
@@ -44,13 +44,40 @@ func ScreenshotScale(sel interface{}, scale float64, picbuf *[]byte, opts ...Que
 		if len(nodes) < 1 {
 			return fmt.Errorf("selector %q did not return any nodes", sel)
 		}
+		return ScreenshotNodes(nodes, scale, picbuf).Do(ctx)
+	}, append(opts, NodeVisible)...)
+}
 
-		// get box model
+// ScreenshotNodes is an action that captures/takes a screenshot of the
+// specified nodes, by calculating the extents of the top most left node and
+// bottom most right node.
+func ScreenshotNodes(nodes []*cdp.Node, scale float64, picbuf *[]byte) Action {
+	if len(nodes) == 0 {
+		panic("nodes must be non-empty")
+	}
+	if picbuf == nil {
+		panic("picbuf cannot be nil")
+	}
+
+	return ActionFunc(func(ctx context.Context) error {
 		var clip page.Viewport
+
+		// get box model of first node
 		if err := callFunctionOnNode(ctx, nodes[0], getClientRectJS, &clip); err != nil {
 			return err
 		}
 
+		// remainder
+		for _, node := range nodes[1:] {
+			var v page.Viewport
+			// get box model of first node
+			if err := callFunctionOnNode(ctx, node, getClientRectJS, &v); err != nil {
+				return err
+			}
+			clip.X, clip.Width = extents(clip.X, clip.Width, v.X, v.Width)
+			clip.Y, clip.Height = extents(clip.Y, clip.Height, v.Y, v.Height)
+		}
+
 		// The "Capture node screenshot" command does not handle fractional dimensions properly.
 		// Let's align with puppeteer:
 		// https://github.com/puppeteer/puppeteer/blob/bba3f41286908ced8f03faf98242d4c3359a5efc/src/common/Page.ts#L2002-L2011
@@ -73,7 +100,7 @@ func ScreenshotScale(sel interface{}, scale float64, picbuf *[]byte, opts ...Que
 
 		*picbuf = buf
 		return nil
-	}, append(opts, NodeVisible)...)
+	})
 }
 
 // CaptureScreenshot is an action that captures/takes a screenshot of the
@@ -133,3 +160,23 @@ func FullScreenshot(res *[]byte, quality int) EmulateAction {
 		return nil
 	})
 }
+
+func extents(m, n, o, p float64) (float64, float64) {
+	a := min(m, o)
+	b := max(m+n, o+p)
+	return a, b - a
+}
+
+func min(a, b float64) float64 {
+	if a < b {
+		return a
+	}
+	return b
+}
+
+func max(a, b float64) float64 {
+	if a > b {
+		return a
+	}
+	return b
+}