refactor: custom query handlers and global bindings (#9678)

packages/puppeteer-core/src/common/AriaQueryHandler.ts

@@ -19,8 +19,6 @@ import {Protocol} from 'devtools-protocol';
 import {ElementHandle} from '../api/ElementHandle.js';
 import {assert} from '../util/assert.js';
 import {CDPSession} from './Connection.js';
-import type {Frame} from './Frame.js';
-import type {WaitForSelectorOptions} from './IsolatedWorld.js';
 import {IterableUtil} from './IterableUtil.js';
 import {QueryHandler, QuerySelector} from './QueryHandler.js';
 import {AwaitableIterable} from './types.js';
@@ -87,20 +85,16 @@ const parseARIASelector = (selector: string): ARIASelector => {
   return queryOptions;
 };
 
-/**
- * @internal
- */
-export interface ARIAQuerySelectorContext {
-  __ariaQuerySelector(node: Node, selector: string): Promise<Node | null>;
-}
-
 /**
  * @internal
  */
 export class ARIAQueryHandler extends QueryHandler {
-  static override querySelector: QuerySelector = async (node, selector) => {
-    const context = globalThis as unknown as ARIAQuerySelectorContext;
-    return context.__ariaQuerySelector(node, selector);
+  static override querySelector: QuerySelector = async (
+    node,
+    selector,
+    {ariaQuerySelector}
+  ) => {
+    return ariaQuerySelector(node, selector);
   };
 
   static override async *queryAll(
@@ -124,17 +118,4 @@ export class ARIAQueryHandler extends QueryHandler {
   ): Promise<ElementHandle<Node> | null> => {
     return (await IterableUtil.first(this.queryAll(element, selector))) ?? null;
   };
-
-  static override async waitFor(
-    elementOrFrame: ElementHandle<Node> | Frame,
-    selector: string,
-    options: WaitForSelectorOptions
-  ): Promise<ElementHandle<Node> | null> {
-    return super.waitFor(
-      elementOrFrame,
-      selector,
-      options,
-      new Map([['__ariaQuerySelector', this.queryOne]])
-    );
-  }
 }

packages/puppeteer-core/src/common/Binding.ts

@@ -6,16 +6,19 @@ import {debugError} from './util.js';
 /**
  * @internal
  */
-export class Binding {
+export class Binding<Params extends unknown[] = any[]> {
   #name: string;
-  #fn: (...args: unknown[]) => unknown;
-  constructor(name: string, fn: (...args: unknown[]) => unknown) {
+  #fn: (...args: Params) => unknown;
+  constructor(name: string, fn: (...args: Params) => unknown) {
     this.#name = name;
     this.#fn = fn;
   }
 
+  get name(): string {
+    return this.#name;
+  }
+
   /**
-   *
    * @param context - Context to run the binding in; the context should have
    * the binding added to it beforehand.
    * @param id - ID of the call. This should come from the CDP
@@ -25,7 +28,7 @@ export class Binding {
   async run(
     context: ExecutionContext,
     id: number,
-    args: unknown[],
+    args: Params,
     isTrivial: boolean
   ): Promise<void> {
     const garbage = [];

packages/puppeteer-core/src/common/CustomQueryHandler.ts

@@ -14,13 +14,11 @@
  * limitations under the License.
  */
 
-import {QueryHandler} from './QueryHandler.js';
-import {getQueryHandlerByName} from './GetQueryHandler.js';
-
-/**
- * @internal
- */
-export const customQueryHandlers = new Map<string, typeof QueryHandler>();
+import type PuppeteerUtil from '../injected/injected.js';
+import {assert} from '../util/assert.js';
+import {interpolateFunction, stringifyFunction} from '../util/Function.js';
+import {QueryHandler, QuerySelector, QuerySelectorAll} from './QueryHandler.js';
+import {scriptInjector} from './ScriptInjector.js';
 
 /**
  * @public
@@ -36,6 +34,154 @@ export interface CustomQueryHandler {
   queryAll?: (node: Node, selector: string) => Iterable<Node>;
 }
 
+/**
+ * The registry of {@link CustomQueryHandler | custom query handlers}.
+ *
+ * @example
+ *
+ * ```ts
+ * Puppeteer.customQueryHandlers.register('lit', { … });
+ * const aHandle = await page.$('lit/…');
+ * ```
+ *
+ * @internal
+ */
+export class CustomQueryHandlerRegistry {
+  #handlers = new Map<
+    string,
+    [registerScript: string, Handler: typeof QueryHandler]
+  >();
+
+  /**
+   * @internal
+   */
+  get(name: string): typeof QueryHandler | undefined {
+    const handler = this.#handlers.get(name);
+    return handler ? handler[1] : undefined;
+  }
+
+  /**
+   * Registers a {@link CustomQueryHandler | custom query handler}.
+   *
+   * @remarks
+   * After registration, the handler can be used everywhere where a selector is
+   * expected by prepending the selection string with `<name>/`. The name is
+   * only allowed to consist of lower- and upper case latin letters.
+   *
+   * @example
+   *
+   * ```ts
+   * Puppeteer.customQueryHandlers.register('lit', { … });
+   * const aHandle = await page.$('lit/…');
+   * ```
+   *
+   * @param name - Name to register under.
+   * @param queryHandler - {@link CustomQueryHandler | Custom query handler} to
+   * register.
+   *
+   * @internal
+   */
+  register(name: string, handler: CustomQueryHandler): void {
+    if (this.#handlers.has(name)) {
+      throw new Error(`Cannot register over existing handler: ${name}`);
+    }
+    assert(
+      !this.#handlers.has(name),
+      `Cannot register over existing handler: ${name}`
+    );
+    assert(
+      /^[a-zA-Z]+$/.test(name),
+      `Custom query handler names may only contain [a-zA-Z]`
+    );
+    assert(
+      handler.queryAll || handler.queryOne,
+      `At least one query method must be implemented.`
+    );
+
+    const Handler = class extends QueryHandler {
+      static override querySelectorAll: QuerySelectorAll = interpolateFunction(
+        (node, selector, PuppeteerUtil) => {
+          return PuppeteerUtil.customQuerySelectors
+            .get(PLACEHOLDER('name'))!
+            .querySelectorAll(node, selector);
+        },
+        {name: JSON.stringify(name)}
+      );
+      static override querySelector: QuerySelector = interpolateFunction(
+        (node, selector, PuppeteerUtil) => {
+          return PuppeteerUtil.customQuerySelectors
+            .get(PLACEHOLDER('name'))!
+            .querySelector(node, selector);
+        },
+        {name: JSON.stringify(name)}
+      );
+    };
+    const registerScript = interpolateFunction(
+      (PuppeteerUtil: PuppeteerUtil) => {
+        PuppeteerUtil.customQuerySelectors.register(PLACEHOLDER('name'), {
+          queryAll: PLACEHOLDER('queryAll'),
+          queryOne: PLACEHOLDER('queryOne'),
+        });
+      },
+      {
+        name: JSON.stringify(name),
+        queryAll: handler.queryAll
+          ? stringifyFunction(handler.queryAll)
+          : String(undefined),
+        queryOne: handler.queryOne
+          ? stringifyFunction(handler.queryOne)
+          : String(undefined),
+      }
+    ).toString();
+
+    this.#handlers.set(name, [registerScript, Handler]);
+    scriptInjector.append(registerScript);
+  }
+
+  /**
+   * Unregisters the {@link CustomQueryHandler | custom query handler} for the
+   * given name.
+   *
+   * @throws `Error` if there is no handler under the given name.
+   *
+   * @internal
+   */
+  unregister(name: string): void {
+    const handler = this.#handlers.get(name);
+    if (!handler) {
+      throw new Error(`Cannot unregister unknown handler: ${name}`);
+    }
+    scriptInjector.pop(handler[0]);
+    this.#handlers.delete(name);
+  }
+
+  /**
+   * Gets the names of all {@link CustomQueryHandler | custom query handlers}.
+   *
+   * @internal
+   */
+  names(): string[] {
+    return [...this.#handlers.keys()];
+  }
+
+  /**
+   * Unregisters all custom query handlers.
+   *
+   * @internal
+   */
+  clear(): void {
+    for (const [registerScript] of this.#handlers) {
+      scriptInjector.pop(registerScript);
+    }
+    this.#handlers.clear();
+  }
+}
+
+/**
+ * @internal
+ */
+export const customQueryHandlers = new CustomQueryHandlerRegistry();
+
 /**
  * @deprecated Import {@link Puppeteer} and use the static method
  * {@link Puppeteer.registerCustomQueryHandler}
@@ -46,22 +192,7 @@ export function registerCustomQueryHandler(
   name: string,
   handler: CustomQueryHandler
 ): void {
-  if (getQueryHandlerByName(name)) {
-    throw new Error(`A query handler named "${name}" already exists`);
-  }
-
-  const isValidName = /^[a-zA-Z]+$/.test(name);
-  if (!isValidName) {
-    throw new Error(`Custom query handler names may only contain [a-zA-Z]`);
-  }
-
-  customQueryHandlers.set(
-    name,
-    class extends QueryHandler {
-      static override querySelector = handler.queryOne;
-      static override querySelectorAll = handler.queryAll;
-    }
-  );
+  customQueryHandlers.register(name, handler);
 }
 
 /**
@@ -71,7 +202,7 @@ export function registerCustomQueryHandler(
  * @public
  */
 export function unregisterCustomQueryHandler(name: string): void {
-  customQueryHandlers.delete(name);
+  customQueryHandlers.unregister(name);
 }
 
 /**
@@ -81,7 +212,7 @@ export function unregisterCustomQueryHandler(name: string): void {
  * @public
  */
 export function customQueryHandlerNames(): string[] {
-  return [...customQueryHandlers.keys()];
+  return customQueryHandlers.names();
 }
 
 /**

packages/puppeteer-core/src/common/ElementHandle.ts

@@ -176,9 +176,9 @@ export class CDPElementHandle<
   override async $<Selector extends string>(
     selector: Selector
   ): Promise<CDPElementHandle<NodeFor<Selector>> | null> {
-    const {updatedSelector, queryHandler} =
+    const {updatedSelector, QueryHandler} =
       getQueryHandlerAndSelector(selector);
-    return (await queryHandler.queryOne(
+    return (await QueryHandler.queryOne(
       this,
       updatedSelector
     )) as CDPElementHandle<NodeFor<Selector>> | null;
@@ -187,10 +187,10 @@ export class CDPElementHandle<
   override async $$<Selector extends string>(
     selector: Selector
   ): Promise<Array<CDPElementHandle<NodeFor<Selector>>>> {
-    const {updatedSelector, queryHandler} =
+    const {updatedSelector, QueryHandler} =
       getQueryHandlerAndSelector(selector);
     return IterableUtil.collect(
-      queryHandler.queryAll(this, updatedSelector)
+      QueryHandler.queryAll(this, updatedSelector)
     ) as Promise<Array<CDPElementHandle<NodeFor<Selector>>>>;
   }
 
@@ -256,9 +256,9 @@ export class CDPElementHandle<
     selector: Selector,
     options: WaitForSelectorOptions = {}
   ): Promise<CDPElementHandle<NodeFor<Selector>> | null> {
-    const {updatedSelector, queryHandler} =
+    const {updatedSelector, QueryHandler} =
       getQueryHandlerAndSelector(selector);
-    return (await queryHandler.waitFor(
+    return (await QueryHandler.waitFor(
       this,
       updatedSelector,
       options

packages/puppeteer-core/src/common/ExecutionContext.ts

@@ -15,18 +15,20 @@
  */
 
 import {Protocol} from 'devtools-protocol';
-import {source as injectedSource} from '../generated/injected.js';
+import {JSHandle} from '../api/JSHandle.js';
 import type PuppeteerUtil from '../injected/injected.js';
+import {stringifyFunction} from '../util/Function.js';
+import {ARIAQueryHandler} from './AriaQueryHandler.js';
+import {Binding} from './Binding.js';
 import {CDPSession} from './Connection.js';
 import {IsolatedWorld} from './IsolatedWorld.js';
-import {JSHandle} from '../api/JSHandle.js';
 import {LazyArg} from './LazyArg.js';
+import {scriptInjector} from './ScriptInjector.js';
 import {EvaluateFunc, HandleFor} from './types.js';
 import {
   createJSHandle,
   getExceptionMessage,
   isString,
-  stringifyFunction,
   valueFromRemoteObject,
 } from './util.js';
 import {CDPJSHandle} from './JSHandle.js';
@@ -94,16 +96,35 @@ export class ExecutionContext {
 
   #puppeteerUtil?: Promise<JSHandle<PuppeteerUtil>>;
   get puppeteerUtil(): Promise<JSHandle<PuppeteerUtil>> {
-    if (!this.#puppeteerUtil) {
-      this.#puppeteerUtil = this.evaluateHandle(
-        `(() => {
-            const module = {};
-            ${injectedSource}
-            return module.exports.default;
-          })()`
-      ) as Promise<JSHandle<PuppeteerUtil>>;
+    scriptInjector.inject(script => {
+      if (this.#puppeteerUtil) {
+        this.#puppeteerUtil.then(handle => {
+          handle.dispose();
+        });
+      }
+      this.#puppeteerUtil = Promise.all([
+        this.#installGlobalBinding(
+          new Binding('__ariaQuerySelector', ARIAQueryHandler.queryOne)
+        ),
+        this.evaluateHandle(script) as Promise<JSHandle<PuppeteerUtil>>,
+      ]).then(([, util]) => {
+        return util;
+      });
+    }, !this.#puppeteerUtil);
+    return this.#puppeteerUtil as Promise<JSHandle<PuppeteerUtil>>;
+  }
+
+  async #installGlobalBinding(binding: Binding<any[]>) {
+    try {
+      if (this._world) {
+        this._world._bindings.set(binding.name, binding);
+        await this._world._addBindingToContext(this, binding.name);
+      }
+    } catch {
+      // If the binding cannot be added, then either the browser doesn't support
+      // bindings (e.g. Firefox) or the context is broken. Either breakage is
+      // okay, so we ignore the error.
     }
-    return this.#puppeteerUtil;
   }
 
   /**
@@ -272,8 +293,7 @@ export class ExecutionContext {
     let callFunctionOnPromise;
     try {
       callFunctionOnPromise = this._client.send('Runtime.callFunctionOn', {
-        functionDeclaration:
-          stringifyFunction(pageFunction) + '\n' + suffix + '\n',
+        functionDeclaration: `${stringifyFunction(pageFunction)}\n${suffix}\n`,
         executionContextId: this._contextId,
         arguments: await Promise.all(args.map(convertArgument.bind(this))),
         returnByValue,

packages/puppeteer-core/src/common/Frame.ts

@@ -617,9 +617,9 @@ export class Frame {
     selector: Selector,
     options: WaitForSelectorOptions = {}
   ): Promise<ElementHandle<NodeFor<Selector>> | null> {
-    const {updatedSelector, queryHandler} =
+    const {updatedSelector, QueryHandler} =
       getQueryHandlerAndSelector(selector);
-    return (await queryHandler.waitFor(
+    return (await QueryHandler.waitFor(
       this,
       updatedSelector,
       options

packages/puppeteer-core/src/common/GetQueryHandler.ts

@@ -48,21 +48,23 @@ export function getQueryHandlerByName(
  */
 export function getQueryHandlerAndSelector(selector: string): {
   updatedSelector: string;
-  queryHandler: typeof QueryHandler;
+  QueryHandler: typeof QueryHandler;
 } {
   for (const handlerMap of [
-    customQueryHandlers,
+    customQueryHandlers.names().map(name => {
+      return [name, customQueryHandlers.get(name)!] as const;
+    }),
     Object.entries(BUILTIN_QUERY_HANDLERS),
   ]) {
-    for (const [name, queryHandler] of handlerMap) {
+    for (const [name, QueryHandler] of handlerMap) {
       for (const separator of QUERY_SEPARATORS) {
         const prefix = `${name}${separator}`;
         if (selector.startsWith(prefix)) {
           selector = selector.slice(prefix.length);
-          return {updatedSelector: selector, queryHandler};
+          return {updatedSelector: selector, QueryHandler};
         }
       }
     }
   }
-  return {updatedSelector: selector, queryHandler: CSSQueryHandler};
+  return {updatedSelector: selector, QueryHandler: CSSQueryHandler};
 }

packages/puppeteer-core/src/common/IsolatedWorld.ts

@@ -41,6 +41,7 @@ import {TaskManager, WaitTask} from './WaitTask.js';
 import type {ElementHandle} from '../api/ElementHandle.js';
 import {Binding} from './Binding.js';
 import {LazyArg} from './LazyArg.js';
+import {stringifyFunction} from '../util/Function.js';
 
 /**
  * @public
@@ -455,7 +456,7 @@ export class IsolatedWorld {
         LazyArg.create(context => {
           return context.puppeteerUtil;
         }),
-        queryOne.toString(),
+        stringifyFunction(queryOne as (...args: unknown[]) => unknown),
         selector,
         root,
         waitForVisible ? true : waitForHidden ? false : undefined

packages/puppeteer-core/src/common/Puppeteer.ts

@@ -20,13 +20,7 @@ import {
   _connectToCDPBrowser,
 } from './BrowserConnector.js';
 import {ConnectionTransport} from './ConnectionTransport.js';
-import {
-  clearCustomQueryHandlers,
-  CustomQueryHandler,
-  customQueryHandlerNames,
-  registerCustomQueryHandler,
-  unregisterCustomQueryHandler,
-} from './CustomQueryHandler.js';
+import {CustomQueryHandler, customQueryHandlers} from './CustomQueryHandler.js';
 
 /**
  * Settings that are common to the Puppeteer class, regardless of environment.
@@ -58,9 +52,18 @@ export interface ConnectOptions extends BrowserConnectOptions {
  * instance of {@link PuppeteerNode} when you import or require `puppeteer`.
  * That class extends `Puppeteer`, so has all the methods documented below as
  * well as all that are defined on {@link PuppeteerNode}.
+ *
  * @public
  */
 export class Puppeteer {
+  /**
+   * Operations for {@link CustomQueryHandler | custom query handlers}. See
+   * {@link CustomQueryHandlerRegistry}.
+   *
+   * @internal
+   */
+  static customQueryHandlers = customQueryHandlers;
+
   /**
    * Registers a {@link CustomQueryHandler | custom query handler}.
    *
@@ -87,28 +90,28 @@ export class Puppeteer {
     name: string,
     queryHandler: CustomQueryHandler
   ): void {
-    return registerCustomQueryHandler(name, queryHandler);
+    return this.customQueryHandlers.register(name, queryHandler);
   }
 
   /**
    * Unregisters a custom query handler for a given name.
    */
   static unregisterCustomQueryHandler(name: string): void {
-    return unregisterCustomQueryHandler(name);
+    return this.customQueryHandlers.unregister(name);
   }
 
   /**
    * Gets the names of all custom query handlers.
    */
   static customQueryHandlerNames(): string[] {
-    return customQueryHandlerNames();
+    return this.customQueryHandlers.names();
   }
 
   /**
    * Unregisters all custom query handlers.
    */
   static clearCustomQueryHandlers(): void {
-    return clearCustomQueryHandlers();
+    return this.customQueryHandlers.clear();
   }
 
   /**

packages/puppeteer-core/src/common/QueryHandler.ts

@@ -17,9 +17,9 @@
 import {ElementHandle} from '../api/ElementHandle.js';
 import type PuppeteerUtil from '../injected/injected.js';
 import {assert} from '../util/assert.js';
-import {createFunction} from '../util/Function.js';
-import {transposeIterableHandle} from './HandleIterator.js';
+import {interpolateFunction, stringifyFunction} from '../util/Function.js';
 import type {Frame} from './Frame.js';
+import {transposeIterableHandle} from './HandleIterator.js';
 import type {WaitForSelectorOptions} from './IsolatedWorld.js';
 import {MAIN_WORLD, PUPPETEER_WORLD} from './IsolatedWorlds.js';
 import {LazyArg} from './LazyArg.js';
@@ -56,28 +56,23 @@ export class QueryHandler {
       return this.querySelector;
     }
     if (!this.querySelectorAll) {
-      throw new Error('Cannot create default query selector');
+      throw new Error('Cannot create default `querySelector`.');
     }
 
-    const querySelector: QuerySelector = async (
-      node,
-      selector,
-      PuppeteerUtil
-    ) => {
-      const querySelectorAll =
-        'FUNCTION_DEFINITION' as unknown as QuerySelectorAll;
-      const results = querySelectorAll(node, selector, PuppeteerUtil);
-      for await (const result of results) {
-        return result;
+    return (this.querySelector = interpolateFunction(
+      async (node, selector, PuppeteerUtil) => {
+        const querySelectorAll: QuerySelectorAll =
+          PLACEHOLDER('querySelectorAll');
+        const results = querySelectorAll(node, selector, PuppeteerUtil);
+        for await (const result of results) {
+          return result;
+        }
+        return null;
+      },
+      {
+        querySelectorAll: stringifyFunction(this.querySelectorAll),
       }
-      return null;
-    };
-
-    return (this.querySelector = createFunction(
-      querySelector
-        .toString()
-        .replace("'FUNCTION_DEFINITION'", this.querySelectorAll.toString())
-    ) as typeof querySelector);
+    ));
   }
 
   static get _querySelectorAll(): QuerySelectorAll {
@@ -85,26 +80,21 @@ export class QueryHandler {
       return this.querySelectorAll;
     }
     if (!this.querySelector) {
-      throw new Error('Cannot create default query selector');
+      throw new Error('Cannot create default `querySelectorAll`.');
     }
 
-    const querySelectorAll: QuerySelectorAll = async function* (
-      node,
-      selector,
-      PuppeteerUtil
-    ) {
-      const querySelector = 'FUNCTION_DEFINITION' as unknown as QuerySelector;
-      const result = await querySelector(node, selector, PuppeteerUtil);
-      if (result) {
-        yield result;
+    return (this.querySelectorAll = interpolateFunction(
+      async function* (node, selector, PuppeteerUtil) {
+        const querySelector: QuerySelector = PLACEHOLDER('querySelector');
+        const result = await querySelector(node, selector, PuppeteerUtil);
+        if (result) {
+          yield result;
+        }
+      },
+      {
+        querySelector: stringifyFunction(this.querySelector),
       }
-    };
-
-    return (this.querySelectorAll = createFunction(
-      querySelectorAll
-        .toString()
-        .replace("'FUNCTION_DEFINITION'", this.querySelector.toString())
-    ) as typeof querySelectorAll);
+    ));
   }
 
   /**
@@ -160,8 +150,7 @@ export class QueryHandler {
   static async waitFor(
     elementOrFrame: ElementHandle<Node> | Frame,
     selector: string,
-    options: WaitForSelectorOptions,
-    bindings = new Map<string, (...args: never[]) => unknown>()
+    options: WaitForSelectorOptions
   ): Promise<ElementHandle<Node> | null> {
     let frame: Frame;
     let element: ElementHandle<Node> | undefined;
@@ -175,8 +164,7 @@ export class QueryHandler {
       this._querySelector,
       element,
       selector,
-      options,
-      bindings
+      options
     );
     if (element) {
       await element.dispose();

packages/puppeteer-core/src/common/ScriptInjector.ts

@@ -0,0 +1,49 @@
+import {source as injectedSource} from '../generated/injected.js';
+
+class ScriptInjector {
+  #updated = false;
+  #amendments = new Set<string>();
+
+  // Appends a statement of the form `(PuppeteerUtil) => {...}`.
+  append(statement: string): void {
+    this.#update(() => {
+      this.#amendments.add(statement);
+    });
+  }
+
+  pop(statement: string): void {
+    this.#update(() => {
+      this.#amendments.delete(statement);
+    });
+  }
+
+  inject(inject: (script: string) => void, force = false) {
+    if (this.#updated || force) {
+      inject(this.#get());
+    }
+    this.#updated = false;
+  }
+
+  #update(callback: () => void): void {
+    callback();
+    this.#updated = true;
+  }
+
+  #get(): string {
+    return `(() => {
+      const module = {};
+      ${injectedSource}
+      ${[...this.#amendments]
+        .map(statement => {
+          return `(${statement})(module.exports.default);`;
+        })
+        .join('')}
+      return module.exports.default;
+    })()`;
+  }
+}
+
+/**
+ * @internal
+ */
+export const scriptInjector = new ScriptInjector();

packages/puppeteer-core/src/common/WaitTask.ts

@@ -18,6 +18,7 @@ import {ElementHandle} from '../api/ElementHandle.js';
 import {JSHandle} from '../api/JSHandle.js';
 import type {Poller} from '../injected/Poller.js';
 import {createDeferredPromise} from '../util/DeferredPromise.js';
+import {stringifyFunction} from '../util/Function.js';
 import {Binding} from './Binding.js';
 import {TimeoutError} from './Errors.js';
 import {IsolatedWorld} from './IsolatedWorld.js';
@@ -68,7 +69,7 @@ export class WaitTask<T = unknown> {
         this.#fn = `() => {return (${fn});}`;
         break;
       default:
-        this.#fn = fn.toString();
+        this.#fn = stringifyFunction(fn);
         break;
     }
     this.#args = args;

packages/puppeteer-core/src/common/bidi/Page.ts

@@ -17,10 +17,11 @@
 import {Page as PageBase} from '../../api/Page.js';
 import {Connection} from './Connection.js';
 import type {EvaluateFunc, HandleFor} from '../types.js';
-import {isString, stringifyFunction} from '../util.js';
+import {isString} from '../util.js';
 import {BidiSerializer} from './Serializer.js';
 import {JSHandle} from './JSHandle.js';
 import {Reference} from './types.js';
+import {stringifyFunction} from '../../util/Function.js';
 
 /**
  * @internal

packages/puppeteer-core/src/common/util.ts

@@ -438,29 +438,3 @@ export async function getReadableFromProtocolStream(
     },
   });
 }
-
-/**
- * @internal
- */
-export function stringifyFunction(expression: Function): string {
-  let functionText = expression.toString();
-  try {
-    new Function('(' + functionText + ')');
-  } catch (error) {
-    // This means we might have a function shorthand. Try another
-    // time prefixing 'function '.
-    if (functionText.startsWith('async ')) {
-      functionText =
-        'async function ' + functionText.substring('async '.length);
-    } else {
-      functionText = 'function ' + functionText;
-    }
-    try {
-      new Function('(' + functionText + ')');
-    } catch (error) {
-      // We tried hard to serialize, but there's a weird beast here.
-      throw new Error('Passed function is not well-serializable!');
-    }
-  }
-  return functionText;
-}

packages/puppeteer-core/src/injected/ARIAQuerySelector.ts

@@ -0,0 +1,31 @@
+/**
+ * Copyright 2022 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+declare global {
+  interface Window {
+    /**
+     * @internal
+     */
+    __ariaQuerySelector(root: Node, selector: string): Promise<Element | null>;
+  }
+}
+
+export const ariaQuerySelector = (
+  root: Node,
+  selector: string
+): Promise<Element | null> => {
+  return window.__ariaQuerySelector(root, selector);
+};

packages/puppeteer-core/src/injected/CustomQuerySelector.ts

@@ -0,0 +1,68 @@
+/**
+ * Copyright 2023 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import type {CustomQueryHandler} from '../common/CustomQueryHandler.js';
+
+export interface CustomQuerySelector {
+  querySelector(root: Node, selector: string): Node | null;
+  querySelectorAll(root: Node, selector: string): Iterable<Node>;
+}
+
+/**
+ * This class mimics the injected {@link CustomQuerySelectorRegistry}.
+ */
+class CustomQuerySelectorRegistry {
+  #selectors = new Map<string, CustomQuerySelector>();
+
+  register(name: string, handler: CustomQueryHandler): void {
+    if (!handler.queryOne && handler.queryAll) {
+      const querySelectorAll = handler.queryAll;
+      handler.queryOne = (node, selector) => {
+        for (const result of querySelectorAll(node, selector)) {
+          return result;
+        }
+        return null;
+      };
+    } else if (handler.queryOne && !handler.queryAll) {
+      const querySelector = handler.queryOne;
+      handler.queryAll = (node, selector) => {
+        const result = querySelector(node, selector);
+        return result ? [result] : [];
+      };
+    } else if (!handler.queryOne || !handler.queryAll) {
+      throw new Error('At least one query method must be defined.');
+    }
+
+    this.#selectors.set(name, {
+      querySelector: handler.queryOne,
+      querySelectorAll: handler.queryAll!,
+    });
+  }
+
+  unregister(name: string): void {
+    this.#selectors.delete(name);
+  }
+
+  get(name: string): CustomQuerySelector | undefined {
+    return this.#selectors.get(name);
+  }
+
+  clear() {
+    this.#selectors.clear();
+  }
+}
+
+export const customQuerySelectors = new CustomQuerySelectorRegistry();

packages/puppeteer-core/src/injected/injected.ts

@@ -16,26 +16,30 @@
 
 import {createDeferredPromise} from '../util/DeferredPromise.js';
 import {createFunction} from '../util/Function.js';
-import {RAFPoller, MutationPoller, IntervalPoller} from './Poller.js';
+import * as ARIAQuerySelector from './ARIAQuerySelector.js';
+import * as CustomQuerySelectors from './CustomQuerySelector.js';
+import * as PierceQuerySelector from './PierceQuerySelector.js';
+import {IntervalPoller, MutationPoller, RAFPoller} from './Poller.js';
 import {
-  isSuitableNodeForTextMatching,
   createTextContent,
+  isSuitableNodeForTextMatching,
 } from './TextContent.js';
 import * as TextQuerySelector from './TextQuerySelector.js';
-import * as XPathQuerySelector from './XPathQuerySelector.js';
-import * as PierceQuerySelector from './PierceQuerySelector.js';
 import * as util from './util.js';
+import * as XPathQuerySelector from './XPathQuerySelector.js';
 
 /**
  * @internal
  */
 const PuppeteerUtil = Object.freeze({
-  ...util,
-  ...TextQuerySelector,
-  ...XPathQuerySelector,
+  ...ARIAQuerySelector,
+  ...CustomQuerySelectors,
   ...PierceQuerySelector,
-  createFunction,
+  ...TextQuerySelector,
+  ...util,
+  ...XPathQuerySelector,
   createDeferredPromise,
+  createFunction,
   createTextContent,
   IntervalPoller,
   isSuitableNodeForTextMatching,

packages/puppeteer-core/src/util/Function.ts

@@ -33,3 +33,66 @@ export const createFunction = (
   createdFunctions.set(functionValue, fn);
   return fn;
 };
+
+/**
+ * @internal
+ */
+export function stringifyFunction(fn: (...args: never) => unknown): string {
+  let value = fn.toString();
+  try {
+    new Function(`(${value})`);
+  } catch {
+    // This means we might have a function shorthand (e.g. `test(){}`). Let's
+    // try prefixing.
+    let prefix = 'function ';
+    if (value.startsWith('async ')) {
+      prefix = `async ${prefix}`;
+      value = value.substring('async '.length);
+    }
+    value = `${prefix}${value}`;
+    try {
+      new Function(`(${value})`);
+    } catch {
+      // We tried hard to serialize, but there's a weird beast here.
+      throw new Error('Passed function cannot be serialized!');
+    }
+  }
+  return value;
+}
+
+/**
+ * Replaces `PLACEHOLDER`s with the given replacements.
+ *
+ * All replacements must be valid JS code.
+ *
+ * @example
+ *
+ * ```ts
+ * interpolateFunction(() => PLACEHOLDER('test'), {test: 'void 0'});
+ * // Equivalent to () => void 0
+ * ```
+ *
+ * @internal
+ */
+export const interpolateFunction = <T extends (...args: never[]) => unknown>(
+  fn: T,
+  replacements: Record<string, string>
+): T => {
+  let value = stringifyFunction(fn);
+  for (const [name, jsValue] of Object.entries(replacements)) {
+    value = value.replace(
+      new RegExp(`PLACEHOLDER\\(\\s*(?:'${name}'|"${name}")\\s*\\)`, 'g'),
+      jsValue
+    );
+  }
+  return createFunction(value) as unknown as T;
+};
+
+declare global {
+  /**
+   * Used for interpolation with {@link interpolateFunction}.
+   *
+   * @internal
+   */
+  function PLACEHOLDER<T>(name: string): T;
+}

test/src/elementhandle.spec.ts

@@ -702,6 +702,27 @@ describe('ElementHandle specs', function () {
       });
       expect(txtContents).toBe('textcontent');
     });
+
+    it('should work with function shorthands', async () => {
+      const {page} = getTestState();
+      await page.setContent('<div id="not-foo"></div><div id="foo"></div>');
+
+      Puppeteer.registerCustomQueryHandler('getById', {
+        // This is a function shorthand
+        queryOne(_element, selector) {
+          return document.querySelector(`[id="${selector}"]`);
+        },
+      });
+
+      const element = (await page.$(
+        'getById/foo'
+      )) as ElementHandle<HTMLDivElement>;
+      expect(
+        await page.evaluate(element => {
+          return element.id;
+        }, element)
+      ).toBe('foo');
+    });
   });
 
   describe('Element.toElement', () => {