puppeteer-core/src/common/CustomQueryHandler.ts

/**
 * @license
 * Copyright 2023 Google Inc.
 * SPDX-License-Identifier: Apache-2.0
 */

import type PuppeteerUtil from '../injected/injected.js';
import {assert} from '../util/assert.js';
import {interpolateFunction, stringifyFunction} from '../util/Function.js';

import {
  QueryHandler,
  type QuerySelector,
  type QuerySelectorAll,
} from './QueryHandler.js';
import {scriptInjector} from './ScriptInjector.js';

/**
 * @public
 */
export interface CustomQueryHandler {
  /**
   * Searches for a {@link https://developer.mozilla.org/en-US/docs/Web/API/Node | Node} matching the given `selector` from {@link https://developer.mozilla.org/en-US/docs/Web/API/Node | node}.
   */
  queryOne?: (node: Node, selector: string) => Node | null;
  /**
   * Searches for some {@link https://developer.mozilla.org/en-US/docs/Web/API/Node | Nodes} matching the given `selector` from {@link https://developer.mozilla.org/en-US/docs/Web/API/Node | node}.
   */
  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]
  >();

  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.
   */
  register(name: string, handler: CustomQueryHandler): void {
    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.
   */
  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}.
   */
  names(): string[] {
    return [...this.#handlers.keys()];
  }

  /**
   * Unregisters all custom query handlers.
   */
  clear(): void {
    for (const [registerScript] of this.#handlers) {
      scriptInjector.pop(registerScript);
    }
    this.#handlers.clear();
  }
}

/**
 * @internal
 */
export const customQueryHandlers = new CustomQueryHandlerRegistry();