feat: support mask option in screenshot methods (#12072)

Fixes https://github.com/microsoft/playwright/issues/10162

docs/src/api/params.md

@@ -913,11 +913,18 @@ saved to the disk.
 
 Specify screenshot type, defaults to `png`.
 
+## screenshot-option-mask
+- `mask` <[Array]<[Locator]>>
+
+Specify locators that should be masked when the screenshot is taken. Masked elements will be overlayed with
+a pink box `#FF00FF` that completely covers its bounding box.
+
 ## screenshot-options-common-list
 - %%-screenshot-option-disable-animations-%%
 - %%-screenshot-option-omit-background-%%
 - %%-screenshot-option-quality-%%
 - %%-screenshot-option-path-%%
 - %%-screenshot-option-type-%%
+- %%-screenshot-option-mask-%%
 - %%-input-timeout-%%
 

packages/playwright-core/src/client/elementHandle.ts

@@ -16,6 +16,7 @@
 
 import * as channels from '../protocol/channels';
 import { Frame } from './frame';
+import { Locator } from './locator';
 import { JSHandle, serializeArgument, parseResult } from './jsHandle';
 import { ChannelOwner } from './channelOwner';
 import { SelectOption, FilePayload, Rect, SelectOptionOptions } from './types';
@@ -173,10 +174,16 @@ export class ElementHandle<T extends Node = Node> extends JSHandle<T> implements
     return value === undefined ? null : value;
   }
 
-  async screenshot(options: channels.ElementHandleScreenshotOptions & { path?: string } = {}): Promise<Buffer> {
-    const copy = { ...options };
+  async screenshot(options: Omit<channels.ElementHandleScreenshotOptions, 'mask'> & { path?: string, mask?: Locator[] } = {}): Promise<Buffer> {
+    const copy: channels.ElementHandleScreenshotOptions = { ...options, mask: undefined };
     if (!copy.type)
       copy.type = determineScreenshotType(options);
+    if (options.mask) {
+      copy.mask = options.mask.map(locator => ({
+        frame: locator._frame._channel,
+        selector: locator._selector,
+      }));
+    }
     const result = await this._elementChannel.screenshot(copy);
     const buffer = Buffer.from(result.binary, 'base64');
     if (options.path) {

packages/playwright-core/src/client/locator.ts

@@ -26,8 +26,8 @@ import { parseResult, serializeArgument } from './jsHandle';
 import { escapeWithQuotes } from '../utils/stringUtils';
 
 export class Locator implements api.Locator {
-  private _frame: Frame;
-  private _selector: string;
+  _frame: Frame;
+  _selector: string;
 
   constructor(frame: Frame, selector: string, options?: { hasText?: string | RegExp, has?: Locator }) {
     this._frame = frame;
@@ -200,7 +200,7 @@ export class Locator implements api.Locator {
     return this._frame.press(this._selector, key, { strict: true, ...options });
   }
 
-  async screenshot(options: channels.ElementHandleScreenshotOptions & { path?: string } = {}): Promise<Buffer> {
+  async screenshot(options: Omit<channels.ElementHandleScreenshotOptions, 'mask'> & { path?: string, mask?: Locator[] } = {}): Promise<Buffer> {
     return this._withElement((h, timeout) => h.screenshot({ ...options, timeout }), options.timeout);
   }
 

packages/playwright-core/src/client/page.ts

@@ -457,10 +457,16 @@ export class Page extends ChannelOwner<channels.PageChannel> implements api.Page
     await this._channel.setNetworkInterceptionEnabled({ enabled: false });
   }
 
-  async screenshot(options: channels.PageScreenshotOptions & { path?: string } = {}): Promise<Buffer> {
-    const copy = { ...options };
+  async screenshot(options: Omit<channels.PageScreenshotOptions, 'mask'> & { path?: string, mask?: Locator[] } = {}): Promise<Buffer> {
+    const copy: channels.PageScreenshotOptions = { ...options, mask: undefined };
     if (!copy.type)
       copy.type = determineScreenshotType(options);
+    if (options.mask) {
+      copy.mask = options.mask.map(locator => ({
+        frame: locator._frame._channel,
+        selector: locator._selector,
+      }));
+    }
     const result = await this._channel.screenshot(copy);
     const buffer = Buffer.from(result.binary, 'base64');
     if (options.path) {

packages/playwright-core/src/dispatchers/elementHandlerDispatcher.ts

@@ -15,6 +15,7 @@
  */
 
 import { ElementHandle } from '../server/dom';
+import { Frame } from '../server/frames';
 import * as js from '../server/javascript';
 import * as channels from '../protocol/channels';
 import { DispatcherScope, existingDispatcher, lookupNullableDispatcher } from './dispatcher';
@@ -171,7 +172,11 @@ export class ElementHandleDispatcher extends JSHandleDispatcher implements chann
   }
 
   async screenshot(params: channels.ElementHandleScreenshotParams, metadata: CallMetadata): Promise<channels.ElementHandleScreenshotResult> {
-    return { binary: (await this._elementHandle.screenshot(metadata, params)).toString('base64') };
+    const mask: { frame: Frame, selector: string }[] = (params.mask || []).map(({ frame, selector }) => ({
+      frame: (frame as FrameDispatcher)._object,
+      selector,
+    }));
+    return { binary: (await this._elementHandle.screenshot(metadata, { ...params, mask })).toString('base64') };
   }
 
   async querySelector(params: channels.ElementHandleQuerySelectorParams, metadata: CallMetadata): Promise<channels.ElementHandleQuerySelectorResult> {

packages/playwright-core/src/dispatchers/pageDispatcher.ts

@@ -151,7 +151,11 @@ export class PageDispatcher extends Dispatcher<Page, channels.PageChannel> imple
   }
 
   async screenshot(params: channels.PageScreenshotParams, metadata: CallMetadata): Promise<channels.PageScreenshotResult> {
-    return { binary: (await this._page.screenshot(metadata, params)).toString('base64') };
+    const mask: { frame: Frame, selector: string }[] = (params.mask || []).map(({ frame, selector }) => ({
+      frame: (frame as FrameDispatcher)._object,
+      selector,
+    }));
+    return { binary: (await this._page.screenshot(metadata, { ...params, mask })).toString('base64') };
   }
 
   async close(params: channels.PageCloseParams, metadata: CallMetadata): Promise<void> {

packages/playwright-core/src/protocol/channels.ts

@@ -1494,6 +1494,10 @@ export type PageScreenshotParams = {
   fullPage?: boolean,
   disableAnimations?: boolean,
   clip?: Rect,
+  mask?: {
+    frame: FrameChannel,
+    selector: string,
+  }[],
 };
 export type PageScreenshotOptions = {
   timeout?: number,
@@ -1503,6 +1507,10 @@ export type PageScreenshotOptions = {
   fullPage?: boolean,
   disableAnimations?: boolean,
   clip?: Rect,
+  mask?: {
+    frame: FrameChannel,
+    selector: string,
+  }[],
 };
 export type PageScreenshotResult = {
   binary: Binary,
@@ -2798,6 +2806,10 @@ export type ElementHandleScreenshotParams = {
   quality?: number,
   omitBackground?: boolean,
   disableAnimations?: boolean,
+  mask?: {
+    frame: FrameChannel,
+    selector: string,
+  }[],
 };
 export type ElementHandleScreenshotOptions = {
   timeout?: number,
@@ -2805,6 +2817,10 @@ export type ElementHandleScreenshotOptions = {
   quality?: number,
   omitBackground?: boolean,
   disableAnimations?: boolean,
+  mask?: {
+    frame: FrameChannel,
+    selector: string,
+  }[],
 };
 export type ElementHandleScreenshotResult = {
   binary: Binary,

packages/playwright-core/src/protocol/protocol.yml

@@ -1003,6 +1003,13 @@ Page:
         fullPage: boolean?
         disableAnimations: boolean?
         clip: Rect?
+        mask:
+          type: array?
+          items:
+            type: object
+            properties:
+              frame: Frame
+              selector: string
       returns:
         binary: binary
 
@@ -2155,6 +2162,13 @@ ElementHandle:
         quality: number?
         omitBackground: boolean?
         disableAnimations: boolean?
+        mask:
+          type: array?
+          items:
+            type: object
+            properties:
+              frame: Frame
+              selector: string
       returns:
         binary: binary
 

packages/playwright-core/src/protocol/validator.ts

@@ -548,6 +548,10 @@ export function createScheme(tChannel: (name: string) => Validator): Scheme {
     fullPage: tOptional(tBoolean),
     disableAnimations: tOptional(tBoolean),
     clip: tOptional(tType('Rect')),
+    mask: tOptional(tArray(tObject({
+      frame: tChannel('Frame'),
+      selector: tString,
+    }))),
   });
   scheme.PageSetExtraHTTPHeadersParams = tObject({
     headers: tArray(tType('NameValue')),
@@ -1038,6 +1042,10 @@ export function createScheme(tChannel: (name: string) => Validator): Scheme {
     quality: tOptional(tNumber),
     omitBackground: tOptional(tBoolean),
     disableAnimations: tOptional(tBoolean),
+    mask: tOptional(tArray(tObject({
+      frame: tChannel('Frame'),
+      selector: tString,
+    }))),
   });
   scheme.ElementHandleScrollIntoViewIfNeededParams = tObject({
     timeout: tOptional(tNumber),

packages/playwright-core/src/server/dom.ts

@@ -18,6 +18,7 @@ import * as mime from 'mime';
 import * as injectedScriptSource from '../generated/injectedScriptSource';
 import * as channels from '../protocol/channels';
 import { isSessionClosedError } from './protocolError';
+import { ScreenshotMaskOption } from './screenshotter';
 import * as frames from './frames';
 import type { InjectedScript, InjectedScriptPoll, LogEntry, HitTargetInterceptionResult } from './injected/injectedScript';
 import { CallMetadata } from './instrumentation';
@@ -772,7 +773,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
     return this._page._delegate.getBoundingBox(this);
   }
 
-  async screenshot(metadata: CallMetadata, options: types.ElementScreenshotOptions = {}): Promise<Buffer> {
+  async screenshot(metadata: CallMetadata, options: types.ElementScreenshotOptions & ScreenshotMaskOption = {}): Promise<Buffer> {
     const controller = new ProgressController(metadata, this);
     return controller.run(
         progress => this._page._screenshotter.screenshotElement(progress, this, options),

packages/playwright-core/src/server/frames.ts

@@ -34,7 +34,7 @@ import { CallMetadata, internalCallMetadata, SdkObject } from './instrumentation
 import type InjectedScript from './injected/injectedScript';
 import type { ElementStateWithoutStable, FrameExpectParams, InjectedScriptPoll, InjectedScriptProgress } from './injected/injectedScript';
 import { isSessionClosedError } from './protocolError';
-import { isInvalidSelectorError, splitSelectorByFrame, stringifySelector } from './common/selectorParser';
+import { isInvalidSelectorError, splitSelectorByFrame, stringifySelector, ParsedSelector } from './common/selectorParser';
 import { SelectorInfo } from './selectors';
 
 type ContextData = {
@@ -786,6 +786,14 @@ export class Frame extends SdkObject {
     return result;
   }
 
+  async maskSelectors(selectors: ParsedSelector[]): Promise<void> {
+    const context = await this._utilityContext();
+    const injectedScript = await context.injectedScript();
+    await injectedScript.evaluate((injected, { parsed }) => {
+      injected.maskSelectors(parsed);
+    }, { parsed: selectors });
+  }
+
   async querySelectorAll(selector: string): Promise<dom.ElementHandle<Element>[]> {
     const pair = await this.resolveFrameForSelectorNoWait(selector, {});
     if (!pair)

packages/playwright-core/src/server/injected/highlight.ts

@@ -175,6 +175,28 @@ export class Highlight {
     }
   }
 
+  maskElements(elements: Element[]) {
+    const boxes = elements.map(e => e.getBoundingClientRect());
+    const pool = this._highlightElements;
+    this._highlightElements = [];
+    for (const box of boxes) {
+      const highlightElement = pool.length ? pool.shift()! : this._createHighlightElement();
+      highlightElement.style.backgroundColor = '#F0F';
+      highlightElement.style.left = box.x + 'px';
+      highlightElement.style.top = box.y + 'px';
+      highlightElement.style.width = box.width + 'px';
+      highlightElement.style.height = box.height + 'px';
+      highlightElement.style.display = 'block';
+      this._highlightElements.push(highlightElement);
+    }
+
+    for (const highlightElement of pool) {
+      highlightElement.style.display = 'none';
+      this._highlightElements.push(highlightElement);
+    }
+  }
+
+
   private _createHighlightElement(): HTMLElement {
     const highlightElement = document.createElement('x-pw-highlight');
     highlightElement.style.position = 'absolute';

packages/playwright-core/src/server/injected/injectedScript.ts

@@ -873,6 +873,17 @@ export class InjectedScript {
     return error;
   }
 
+  maskSelectors(selectors: ParsedSelector[]) {
+    if (this._highlight)
+      this.hideHighlight();
+    this._highlight = new Highlight(false);
+    this._highlight.install();
+    const elements = [];
+    for (const selector of selectors)
+      elements.push(this.querySelectorAll(selector, document.documentElement));
+    this._highlight.maskElements(elements.flat());
+  }
+
   highlight(selector: ParsedSelector) {
     if (!this._highlight) {
       this._highlight = new Highlight(false);

packages/playwright-core/src/server/page.ts

@@ -20,7 +20,7 @@ import * as frames from './frames';
 import * as input from './input';
 import * as js from './javascript';
 import * as network from './network';
-import { Screenshotter } from './screenshotter';
+import { Screenshotter, ScreenshotMaskOption } from './screenshotter';
 import { TimeoutSettings } from '../utils/timeoutSettings';
 import * as types from './types';
 import { BrowserContext } from './browserContext';
@@ -424,7 +424,7 @@ export class Page extends SdkObject {
     route.continue();
   }
 
-  async screenshot(metadata: CallMetadata, options: types.ScreenshotOptions = {}): Promise<Buffer> {
+  async screenshot(metadata: CallMetadata, options: types.ScreenshotOptions & ScreenshotMaskOption = {}): Promise<Buffer> {
     const controller = new ProgressController(metadata, this);
     return controller.run(
         progress => this._screenshotter.screenshotPage(progress, options),

packages/playwright-core/src/server/screenshotter.ts

@@ -18,9 +18,12 @@
 import * as dom from './dom';
 import { helper } from './helper';
 import { Page } from './page';
+import { Frame } from './frames';
+import { ParsedSelector } from './common/selectorParser';
 import * as types from './types';
 import { Progress } from './progress';
 import { assert } from '../utils/utils';
+import { MultiMap } from '../utils/multimap';
 
 declare global {
   interface Window {
@@ -28,6 +31,10 @@ declare global {
   }
 }
 
+export type ScreenshotMaskOption = {
+  mask?: { frame: Frame, selector: string}[],
+};
+
 export class Screenshotter {
   private _queue = new TaskQueue();
   private _page: Page;
@@ -65,7 +72,7 @@ export class Screenshotter {
     return fullPageSize!;
   }
 
-  async screenshotPage(progress: Progress, options: types.ScreenshotOptions): Promise<Buffer> {
+  async screenshotPage(progress: Progress, options: types.ScreenshotOptions & ScreenshotMaskOption): Promise<Buffer> {
     const format = validateScreenshotOptions(options);
     return this._queue.postTask(async () => {
       const { viewportSize } = await this._originalViewportSize(progress);
@@ -92,7 +99,7 @@ export class Screenshotter {
     });
   }
 
-  async screenshotElement(progress: Progress, handle: dom.ElementHandle, options: types.ElementScreenshotOptions = {}): Promise<Buffer> {
+  async screenshotElement(progress: Progress, handle: dom.ElementHandle, options: types.ElementScreenshotOptions & ScreenshotMaskOption = {}): Promise<Buffer> {
     const format = validateScreenshotOptions(options);
     return this._queue.postTask(async () => {
       const { viewportSize } = await this._originalViewportSize(progress);
@@ -210,7 +217,22 @@ export class Screenshotter {
     }));
   }
 
-  private async _screenshot(progress: Progress, format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, fitsViewport: boolean | undefined, options: types.ElementScreenshotOptions): Promise<Buffer> {
+  async _maskElements(progress: Progress, options: ScreenshotMaskOption) {
+    const framesToParsedSelectors: MultiMap<Frame, ParsedSelector> = new MultiMap();
+    await Promise.all((options.mask || []).map(async ({ frame, selector }) => {
+      const pair = await frame.resolveFrameForSelectorNoWait(selector);
+      if (pair)
+        framesToParsedSelectors.set(pair.frame, pair.info.parsed);
+    }));
+    progress.throwIfAborted(); // Avoid extra work.
+
+    await Promise.all([...framesToParsedSelectors.keys()].map(async frame => {
+      await frame.maskSelectors(framesToParsedSelectors.get(frame));
+    }));
+    progress.cleanupWhenAborted(() => this._page.hideHighlight());
+  }
+
+  private async _screenshot(progress: Progress, format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, fitsViewport: boolean | undefined, options: types.ElementScreenshotOptions & ScreenshotMaskOption): Promise<Buffer> {
     if ((options as any).__testHookBeforeScreenshot)
       await (options as any).__testHookBeforeScreenshot();
     progress.throwIfAborted(); // Screenshotting is expensive - avoid extra work.
@@ -221,8 +243,15 @@ export class Screenshotter {
     }
     progress.throwIfAborted(); // Avoid extra work.
 
+    await this._maskElements(progress, options);
+    progress.throwIfAborted(); // Avoid extra work.
+
     const buffer = await this._page._delegate.takeScreenshot(progress, format, documentRect, viewportRect, options.quality, fitsViewport);
     progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
+
+    await this._page.hideHighlight();
+    progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
+
     if (shouldSetDefaultBackground)
       await this._page._delegate.setBackgroundColor();
     progress.throwIfAborted(); // Avoid side effects.

packages/playwright-core/types/types.d.ts

@@ -8070,6 +8070,12 @@ export interface ElementHandle<T=Node> extends JSHandle<T> {
      */
     disableAnimations?: boolean;
 
+    /**
+     * Specify locators that should be masked when the screenshot is taken. Masked elements will be overlayed with a pink box
+     * `#FF00FF` that completely covers its bounding box.
+     */
+    mask?: Array<Locator>;
+
     /**
      * Hides default white background and allows capturing screenshots with transparency. Not applicable to `jpeg` images.
      * Defaults to `false`.
@@ -9373,6 +9379,12 @@ export interface Locator {
      */
     disableAnimations?: boolean;
 
+    /**
+     * Specify locators that should be masked when the screenshot is taken. Masked elements will be overlayed with a pink box
+     * `#FF00FF` that completely covers its bounding box.
+     */
+    mask?: Array<Locator>;
+
     /**
      * Hides default white background and allows capturing screenshots with transparency. Not applicable to `jpeg` images.
      * Defaults to `false`.
@@ -15722,6 +15734,12 @@ export interface PageScreenshotOptions {
    */
   fullPage?: boolean;
 
+  /**
+   * Specify locators that should be masked when the screenshot is taken. Masked elements will be overlayed with a pink box
+   * `#FF00FF` that completely covers its bounding box.
+   */
+  mask?: Array<Locator>;
+
   /**
    * Hides default white background and allows capturing screenshots with transparency. Not applicable to `jpeg` images.
    * Defaults to `false`.

tests/page/page-screenshot.spec.ts

@@ -16,7 +16,7 @@
  */
 
 import { test as it, expect } from './pageTest';
-import { verifyViewport } from '../config/utils';
+import { verifyViewport, attachFrame } from '../config/utils';
 import path from 'path';
 import fs from 'fs';
 
@@ -338,6 +338,86 @@ it.describe('page screenshot', () => {
       screenshotSeveralTimes()
     ]);
   });
+
+  it.describe('mask option', () => {
+    it('should work', async ({ page, server }) => {
+      await page.setViewportSize({ width: 500, height: 500 });
+      await page.goto(server.PREFIX + '/grid.html');
+      expect(await page.screenshot({
+        mask: [ page.locator('div').nth(5) ],
+      })).toMatchSnapshot('mask-should-work.png');
+    });
+
+    it('should work with locator', async ({ page, server }) => {
+      await page.setViewportSize({ width: 500, height: 500 });
+      await page.goto(server.PREFIX + '/grid.html');
+      const bodyLocator = page.locator('body');
+      expect(await bodyLocator.screenshot({
+        mask: [ page.locator('div').nth(5) ],
+      })).toMatchSnapshot('mask-should-work-with-locator.png');
+    });
+
+    it('should work with elementhandle', async ({ page, server }) => {
+      await page.setViewportSize({ width: 500, height: 500 });
+      await page.goto(server.PREFIX + '/grid.html');
+      const bodyHandle = await page.$('body');
+      expect(await bodyHandle.screenshot({
+        mask: [ page.locator('div').nth(5) ],
+      })).toMatchSnapshot('mask-should-work-with-elementhandle.png');
+    });
+
+    it('should mask multiple elements', async ({ page, server }) => {
+      await page.setViewportSize({ width: 500, height: 500 });
+      await page.goto(server.PREFIX + '/grid.html');
+      expect(await page.screenshot({
+        mask: [
+          page.locator('div').nth(5),
+          page.locator('div').nth(12),
+        ],
+      })).toMatchSnapshot('should-mask-multiple-elements.png');
+    });
+
+    it('should mask inside iframe', async ({ page, server }) => {
+      await page.setViewportSize({ width: 500, height: 500 });
+      await page.goto(server.PREFIX + '/grid.html');
+      await attachFrame(page, 'frame1', server.PREFIX + '/grid.html');
+      await page.addStyleTag({ content: 'iframe { border: none; }' });
+      expect(await page.screenshot({
+        mask: [
+          page.locator('div').nth(5),
+          page.frameLocator('#frame1').locator('div').nth(12),
+        ],
+      })).toMatchSnapshot('should-mask-inside-iframe.png');
+    });
+
+    it('should mask in parallel', async ({ page, server }) => {
+      await page.setViewportSize({ width: 500, height: 500 });
+      await attachFrame(page, 'frame1', server.PREFIX + '/grid.html');
+      await attachFrame(page, 'frame2', server.PREFIX + '/grid.html');
+      await page.addStyleTag({ content: 'iframe { border: none; }' });
+      const screenshots = await Promise.all([
+        page.screenshot({
+          mask: [ page.frameLocator('#frame1').locator('div').nth(1) ],
+        }),
+        page.screenshot({
+          mask: [ page.frameLocator('#frame2').locator('div').nth(3) ],
+        }),
+      ]);
+      expect(screenshots[0]).toMatchSnapshot('should-mask-in-parallel-1.png');
+      expect(screenshots[1]).toMatchSnapshot('should-mask-in-parallel-2.png');
+    });
+
+    it('should remove mask after screenshot', async ({ page, server }) => {
+      await page.setViewportSize({ width: 500, height: 500 });
+      await page.goto(server.PREFIX + '/grid.html');
+      const screenshot1 = await page.screenshot();
+      await page.screenshot({
+        mask: [ page.locator('div').nth(1) ],
+      });
+      const screenshot2 = await page.screenshot();
+      expect(screenshot1.equals(screenshot2)).toBe(true);
+    });
+  });
 });
 
 async function rafraf(page) {
@@ -518,6 +598,7 @@ it.describe('page screenshot animations', () => {
     await page.goto(server.PREFIX + '/rotate-z.html');
     await page.evaluate(async () => {
       window.animation = document.getAnimations()[0];
+      await window.animation.ready;
       window.animation.updatePlaybackRate(0);
       await window.animation.ready;
       window.animation.currentTime = 500;
@@ -615,3 +696,4 @@ it.describe('page screenshot animations', () => {
     ]);
   });
 });
+