From d85935a2cb50dce0508474b8278049e2e5247559 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Janne=20Hy=C3=B6tyl=C3=A4?= <janne.hyoetylae@medartis.com>
Date: Mon, 11 Aug 2025 12:27:29 +0200
Subject: [PATCH] fix: Fix masking in single element screenshots.

Using masks in single element screenshots (`Locator.screenshot()`) never worked, because the mask option failed to be serialized.

Add a serializer for `LocatorImpl`, so that this works everywhere correctly.

Fixes #1790
---
 .../com/microsoft/playwright/impl/PageImpl.java     | 10 ----------
 .../microsoft/playwright/impl/Serialization.java    |  8 ++++++++
 .../microsoft/playwright/TestPageScreenshot.java    | 13 ++++++++++++-
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/playwright/src/main/java/com/microsoft/playwright/impl/PageImpl.java b/playwright/src/main/java/com/microsoft/playwright/impl/PageImpl.java
index 30a8ae573..cf23e4c0c 100644
--- a/playwright/src/main/java/com/microsoft/playwright/impl/PageImpl.java
+++ b/playwright/src/main/java/com/microsoft/playwright/impl/PageImpl.java
@@ -1166,18 +1166,8 @@ private byte[] screenshotImpl(ScreenshotOptions options) {
         }
       }
     }
-    List<Locator> mask = options.mask;
-    options.mask = null;
     JsonObject params = gson().toJsonTree(options).getAsJsonObject();
-    options.mask = mask;
     params.remove("path");
-    if (mask != null) {
-      JsonArray maskArray = new JsonArray();
-      for (Locator locator: mask) {
-        maskArray.add(((LocatorImpl) locator).toProtocol());
-      }
-      params.add("mask", maskArray);
-    }
     JsonObject json = sendMessage("screenshot", params, timeoutSettings.timeout(options.timeout)).getAsJsonObject();
 
     byte[] buffer = Base64.getDecoder().decode(json.get("binary").getAsString());
diff --git a/playwright/src/main/java/com/microsoft/playwright/impl/Serialization.java b/playwright/src/main/java/com/microsoft/playwright/impl/Serialization.java
index 633e75761..8be93f1c1 100644
--- a/playwright/src/main/java/com/microsoft/playwright/impl/Serialization.java
+++ b/playwright/src/main/java/com/microsoft/playwright/impl/Serialization.java
@@ -66,6 +66,7 @@ class Serialization {
     .registerTypeHierarchyAdapter(JSHandleImpl.class, new HandleSerializer())
     .registerTypeAdapter((new TypeToken<Map<String, String>>(){}).getType(), new StringMapSerializer())
     .registerTypeAdapter((new TypeToken<Map<String, Object>>(){}).getType(), new FirefoxUserPrefsSerializer())
+    .registerTypeAdapter(LocatorImpl.class, new LocatorImplSerializer())
     .registerTypeHierarchyAdapter(Path.class, new PathSerializer()).create();
 
   static Gson gson() {
@@ -490,6 +491,13 @@ public JsonElement serialize(E src, Type typeOfSrc, JsonSerializationContext con
     }
   }
 
+  private static class LocatorImplSerializer implements JsonSerializer<LocatorImpl> {
+    @Override
+    public JsonElement serialize(LocatorImpl src, Type typeOfSrc, JsonSerializationContext context) {
+      return src.toProtocol();
+    }
+  }
+
   private static class SameSiteAdapter extends TypeAdapter<SameSiteAttribute> {
     @Override
     public void write(JsonWriter out, SameSiteAttribute value) throws IOException {
diff --git a/playwright/src/test/java/com/microsoft/playwright/TestPageScreenshot.java b/playwright/src/test/java/com/microsoft/playwright/TestPageScreenshot.java
index a756476b3..2e82c9208 100644
--- a/playwright/src/test/java/com/microsoft/playwright/TestPageScreenshot.java
+++ b/playwright/src/test/java/com/microsoft/playwright/TestPageScreenshot.java
@@ -136,7 +136,7 @@ void shouldNotCaptureInfiniteWebAnimations() {
   }
 
   @Test
-  void maskShouldWork() {
+  void maskShouldWorkForPage() {
     page.setViewportSize(500, 500);
     page.navigate(server.PREFIX + "/grid.html");
     byte[] screenshot = page.screenshot(new Page.ScreenshotOptions()
@@ -146,6 +146,17 @@ void maskShouldWork() {
     assertThrows(AssertionFailedError.class, () -> assertArrayEquals(screenshot, originalScreenshot));
   }
 
+  @Test
+  void maskShouldWorkForLocator() {
+    page.navigate(server.PREFIX + "/grid.html");
+    Locator locatorToScreenshot = page.locator("div").first();
+    byte[] screenshot = locatorToScreenshot.screenshot(new Locator.ScreenshotOptions()
+      .setMask(asList(page.locator("img"))));
+    // TODO: toMatchSnapshot is not present in java, so we only checks that masked screenshot is different.
+    byte[] originalScreenshot = locatorToScreenshot.screenshot();
+    assertThrows(AssertionFailedError.class, () -> assertArrayEquals(screenshot, originalScreenshot));
+  }
+
   @Test
   void shouldWorkWithDeviceScaleFactorAndClip() {
     try (BrowserContext context = browser.newContext(new Browser.NewContextOptions()