feat: change QueryHandler to contain QueryOne and QueryAll methods

src/common/DOMWorld.ts

@@ -23,7 +23,7 @@ import { ExecutionContext } from './ExecutionContext';
 import { TimeoutSettings } from './TimeoutSettings';
 import { MouseButton } from './Input';
 import { FrameManager, Frame } from './FrameManager';
-import { getQueryHandlerAndSelector, QueryHandler } from './QueryHandler';
+import { getQueryHandlerAndSelector } from './QueryHandler';
 import {
   SerializableOrJSHandle,
   EvaluateHandleFn,
@@ -33,7 +33,10 @@ import { isNode } from '../environment';
 
 // This predicateQueryHandler is declared here so that TypeScript knows about it
 // when it is used in the predicate function below.
-declare const predicateQueryHandler: QueryHandler;
+declare const predicateQueryHandler: (
+  element: Element | Document,
+  selector: string
+) => Element | Element[] | NodeListOf<Element>;
 
 /**
  * @public
@@ -497,16 +500,13 @@ export class DOMWorld {
     const title = `${isXPath ? 'XPath' : 'selector'} "${selectorOrXPath}"${
       waitForHidden ? ' to be hidden' : ''
     }`;
-    const {
-      updatedSelector,
-      queryHandler,
-    } = getQueryHandlerAndSelector(selectorOrXPath, (element, selector) =>
-      document.querySelector(selector)
+    const { updatedSelector, queryHandler } = getQueryHandlerAndSelector(
+      selectorOrXPath
     );
     const waitTask = new WaitTask(
       this,
       predicate,
-      queryHandler,
+      queryHandler.queryOne,
       title,
       polling,
       timeout,

src/common/JSHandle.ts

@@ -772,14 +772,14 @@ export class ElementHandle<
    * the return value resolves to `null`.
    */
   async $(selector: string): Promise<ElementHandle | null> {
-    const defaultHandler = (element: Element, selector: string) =>
-      element.querySelector(selector);
     const { updatedSelector, queryHandler } = getQueryHandlerAndSelector(
-      selector,
-      defaultHandler
+      selector
     );
 
-    const handle = await this.evaluateHandle(queryHandler, updatedSelector);
+    const handle = await this.evaluateHandle(
+      queryHandler.queryOne,
+      updatedSelector
+    );
     const element = handle.asElement();
     if (element) return element;
     await handle.dispose();
@@ -791,15 +791,12 @@ export class ElementHandle<
    * the return value resolves to `[]`.
    */
   async $$(selector: string): Promise<ElementHandle[]> {
-    const defaultHandler = (element: Element, selector: string) =>
-      element.querySelectorAll(selector);
     const { updatedSelector, queryHandler } = getQueryHandlerAndSelector(
-      selector,
-      defaultHandler
+      selector
     );
 
     const arrayHandle = await this.evaluateHandle(
-      queryHandler,
+      queryHandler.queryAll,
       updatedSelector
     );
     const properties = await arrayHandle.getProperties();
@@ -890,15 +887,17 @@ export class ElementHandle<
     ) => ReturnType | Promise<ReturnType>,
     ...args: SerializableOrJSHandle[]
   ): Promise<WrapElementHandle<ReturnType>> {
-    const defaultHandler = (element: Element, selector: string) =>
-      Array.from(element.querySelectorAll(selector));
     const { updatedSelector, queryHandler } = getQueryHandlerAndSelector(
-      selector,
-      defaultHandler
+      selector
     );
-
+    const queryHandlerToArrayStr = `(element, selector) =>
+      Array.from((${queryHandler.queryAll})(element, selector))`;
+    /* The `eval` ensures that the query handler is inlined
+     * before sending the program text to the browser.
+     */
+    const queryHandlerToArray = eval(queryHandlerToArrayStr);
     const arrayHandle = await this.evaluateHandle(
-      queryHandler,
+      queryHandlerToArray,
       updatedSelector
     );
     const result = await arrayHandle.evaluate<

src/common/QueryHandler.ts

@@ -15,17 +15,18 @@
  */
 
 export interface QueryHandler {
-  (element: Element | Document, selector: string):
-    | Element
-    | Element[]
-    | NodeListOf<Element>;
+  queryOne?: (element: Element | Document, selector: string) => Element;
+  queryAll?: (
+    element: Element | Document,
+    selector: string
+  ) => Element[] | NodeListOf<Element>;
 }
 
 const _customQueryHandlers = new Map<string, QueryHandler>();
 
 export function registerCustomQueryHandler(
   name: string,
-  handler: Function
+  handler: QueryHandler
 ): void {
   if (_customQueryHandlers.get(name))
     throw new Error(`A custom query handler named "${name}" already exists`);
@@ -34,7 +35,7 @@ export function registerCustomQueryHandler(
   if (!isValidName)
     throw new Error(`Custom query handler names may only contain [a-zA-Z]`);
 
-  _customQueryHandlers.set(name, handler as QueryHandler);
+  _customQueryHandlers.set(name, handler);
 }
 
 /**
@@ -53,12 +54,17 @@ export function clearQueryHandlers(): void {
 }
 
 export function getQueryHandlerAndSelector(
-  selector: string,
-  defaultQueryHandler: QueryHandler
+  selector: string
 ): { updatedSelector: string; queryHandler: QueryHandler } {
+  const defaultHandler = {
+    queryOne: (element: Element, selector: string) =>
+      element.querySelector(selector),
+    queryAll: (element: Element, selector: string) =>
+      element.querySelectorAll(selector),
+  };
   const hasCustomQueryHandler = /^[a-zA-Z]+\//.test(selector);
   if (!hasCustomQueryHandler)
-    return { updatedSelector: selector, queryHandler: defaultQueryHandler };
+    return { updatedSelector: selector, queryHandler: defaultHandler };
 
   const index = selector.indexOf('/');
   const name = selector.slice(0, index);

test/elementhandle.spec.ts

@@ -290,10 +290,10 @@ describe('ElementHandle specs', function () {
       await page.setContent('<div id="not-foo"></div><div id="foo"></div>');
 
       // Register.
-      puppeteer.__experimental_registerCustomQueryHandler(
-        'getById',
-        (element, selector) => document.querySelector(`[id="${selector}"]`)
-      );
+      puppeteer.__experimental_registerCustomQueryHandler('getById', {
+        queryOne: (element, selector) =>
+          document.querySelector(`[id="${selector}"]`),
+      });
       const element = await page.$('getById/foo');
       expect(await page.evaluate((element) => element.id, element)).toBe('foo');
 
@@ -332,10 +332,10 @@ describe('ElementHandle specs', function () {
       await page.setContent(
         '<div id="not-foo"></div><div class="foo">Foo1</div><div class="foo baz">Foo2</div>'
       );
-      puppeteer.__experimental_registerCustomQueryHandler(
-        'getByClass',
-        (element, selector) => document.querySelectorAll(`.${selector}`)
-      );
+      puppeteer.__experimental_registerCustomQueryHandler('getByClass', {
+        queryAll: (element, selector) =>
+          document.querySelectorAll(`.${selector}`),
+      });
       const elements = await page.$$('getByClass/foo');
       const classNames = await Promise.all(
         elements.map(
@@ -351,10 +351,10 @@ describe('ElementHandle specs', function () {
       await page.setContent(
         '<div id="not-foo"></div><div class="foo">Foo1</div><div class="foo baz">Foo2</div>'
       );
-      puppeteer.__experimental_registerCustomQueryHandler(
-        'getByClass',
-        (element, selector) => document.querySelectorAll(`.${selector}`)
-      );
+      puppeteer.__experimental_registerCustomQueryHandler('getByClass', {
+        queryAll: (element, selector) =>
+          document.querySelectorAll(`.${selector}`),
+      });
       const elements = await page.$$eval(
         'getByClass/foo',
         (divs) => divs.length
@@ -364,10 +364,9 @@ describe('ElementHandle specs', function () {
     });
     it('should wait correctly with waitForSelector', async () => {
       const { page, puppeteer } = getTestState();
-      puppeteer.__experimental_registerCustomQueryHandler(
-        'getByClass',
-        (element, selector) => element.querySelector(`.${selector}`)
-      );
+      puppeteer.__experimental_registerCustomQueryHandler('getByClass', {
+        queryOne: (element, selector) => element.querySelector(`.${selector}`),
+      });
       const waitFor = page.waitForSelector('getByClass/foo');
 
       // Set the page content after the waitFor has been started.
@@ -380,10 +379,9 @@ describe('ElementHandle specs', function () {
     });
     it('should wait correctly with waitFor', async () => {
       const { page, puppeteer } = getTestState();
-      puppeteer.__experimental_registerCustomQueryHandler(
-        'getByClass',
-        (element, selector) => element.querySelector(`.${selector}`)
-      );
+      puppeteer.__experimental_registerCustomQueryHandler('getByClass', {
+        queryOne: (element, selector) => element.querySelector(`.${selector}`),
+      });
       const waitFor = page.waitFor('getByClass/foo');
 
       // Set the page content after the waitFor has been started.