feat: change QueryHandler to contain QueryOne and QueryAll methods (#6218)

Co-authored-by: Mathias Bynens <mathias@qiwi.be>

src/common/DOMWorld.ts

@@ -26,7 +26,7 @@ import { ExecutionContext } from './ExecutionContext.js';
 import { TimeoutSettings } from './TimeoutSettings.js';
 import { MouseButton } from './Input.js';
 import { FrameManager, Frame } from './FrameManager.js';
-import { getQueryHandlerAndSelector, QueryHandler } from './QueryHandler.js';
+import { getQueryHandlerAndSelector } from './QueryHandler.js';
 import {
   SerializableOrJSHandle,
   EvaluateHandleFn,
@@ -39,7 +39,10 @@ import { isNode } from '../environment.js';
 
 // 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
@@ -506,16 +509,13 @@ export class DOMWorld {
     const title = `${isXPath ? 'XPath' : 'selector'} "${selectorOrXPath}"${
       waitForHidden ? ' to be hidden' : ''
     }`;
-    const {
+    const { updatedSelector, queryHandler } = getQueryHandlerAndSelector(
-      updatedSelector,
+      selectorOrXPath
-      queryHandler,
-    } = getQueryHandlerAndSelector(selectorOrXPath, (element, selector) =>
-      document.querySelector(selector)
     );
     const waitTask = new WaitTask(
       this,
       predicate,
-      queryHandler,
+      queryHandler.queryOne,
       title,
       polling,
       timeout,

src/common/JSHandle.ts

@@ -774,14 +774,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,
+      selector
-      defaultHandler
     );
 
-    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();
@@ -793,19 +793,16 @@ 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,
+      selector
-      defaultHandler
     );
 
-    const arrayHandle = await this.evaluateHandle(
+    const handles = await this.evaluateHandle(
-      queryHandler,
+      queryHandler.queryAll,
       updatedSelector
     );
-    const properties = await arrayHandle.getProperties();
+    const properties = await handles.getProperties();
-    await arrayHandle.dispose();
+    await handles.dispose();
     const result = [];
     for (const property of properties.values()) {
       const elementHandle = property.asElement();
@@ -851,8 +848,8 @@ export class ElementHandle<
     await elementHandle.dispose();
 
     /**
-     * This as is a little unfortunate but helps TS understand the behavour of
+     * This `as` is a little unfortunate but helps TS understand the behavior of
-     * `elementHandle.evaluate`. If evalute returns an element it will return an
+     * `elementHandle.evaluate`. If evaluate returns an element it will return an
      * ElementHandle instance, rather than the plain object. All the
      * WrapElementHandle type does is wrap ReturnType into
      * ElementHandle<ReturnType> if it is an ElementHandle, or leave it alone as
@@ -892,15 +889,16 @@ 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,
+      selector
-      defaultHandler
     );
-
+    const queryHandlerToArray = Function(
+      'element',
+      'selector',
+      `return Array.from((${queryHandler.queryAll})(element, selector));`
+    ) as (...args: unknown[]) => unknown;
     const arrayHandle = await this.evaluateHandle(
-      queryHandler,
+      queryHandlerToArray,
       updatedSelector
     );
     const result = await arrayHandle.evaluate<
@@ -910,8 +908,8 @@ export class ElementHandle<
       ) => ReturnType | Promise<ReturnType>
     >(pageFunction, ...args);
     await arrayHandle.dispose();
-    /* This as exists for the same reason as the `as` in $eval above.
+    /* This `as` exists for the same reason as the `as` in $eval above.
-     * See the comment there for a ful explanation.
+     * See the comment there for a full explanation.
      */
     return result as WrapElementHandle<ReturnType>;
   }

src/common/QueryHandler.ts

@@ -15,17 +15,18 @@
  */
 
 export interface QueryHandler {
-  (element: Element | Document, selector: string):
+  queryOne?: (element: Element | Document, selector: string) => Element | null;
-    | Element
+  queryAll?: (
-    | Element[]
+    element: Element | Document,
-    | NodeListOf<Element>;
+    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,
+  selector: string
-  defaultQueryHandler: QueryHandler
 ): { 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

@@ -293,10 +293,10 @@ describe('ElementHandle specs', function () {
       await page.setContent('<div id="not-foo"></div><div id="foo"></div>');
 
       // Register.
-      puppeteer.__experimental_registerCustomQueryHandler(
+      puppeteer.__experimental_registerCustomQueryHandler('getById', {
-        'getById',
+        queryOne: (element, selector) =>
-        (element, selector) => document.querySelector(`[id="${selector}"]`)
+          document.querySelector(`[id="${selector}"]`),
-      );
+      });
       const element = await page.$('getById/foo');
       expect(
         await page.evaluate<(element: HTMLElement) => string>(
@@ -340,10 +340,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(
+      puppeteer.__experimental_registerCustomQueryHandler('getByClass', {
-        'getByClass',
+        queryAll: (element, selector) =>
-        (element, selector) => document.querySelectorAll(`.${selector}`)
+          document.querySelectorAll(`.${selector}`),
-      );
+      });
       const elements = await page.$$('getByClass/foo');
       const classNames = await Promise.all(
         elements.map(
@@ -362,10 +362,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(
+      puppeteer.__experimental_registerCustomQueryHandler('getByClass', {
-        'getByClass',
+        queryAll: (element, selector) =>
-        (element, selector) => document.querySelectorAll(`.${selector}`)
+          document.querySelectorAll(`.${selector}`),
-      );
+      });
       const elements = await page.$$eval(
         'getByClass/foo',
         (divs) => divs.length
@@ -375,10 +375,9 @@ describe('ElementHandle specs', function () {
     });
     it('should wait correctly with waitForSelector', async () => {
       const { page, puppeteer } = getTestState();
-      puppeteer.__experimental_registerCustomQueryHandler(
+      puppeteer.__experimental_registerCustomQueryHandler('getByClass', {
-        'getByClass',
+        queryOne: (element, selector) => element.querySelector(`.${selector}`),
-        (element, selector) => element.querySelector(`.${selector}`)
+      });
-      );
       const waitFor = page.waitForSelector('getByClass/foo');
 
       // Set the page content after the waitFor has been started.
@@ -391,10 +390,9 @@ describe('ElementHandle specs', function () {
     });
     it('should wait correctly with waitFor', async () => {
       const { page, puppeteer } = getTestState();
-      puppeteer.__experimental_registerCustomQueryHandler(
+      puppeteer.__experimental_registerCustomQueryHandler('getByClass', {
-        'getByClass',
+        queryOne: (element, selector) => element.querySelector(`.${selector}`),
-        (element, selector) => element.querySelector(`.${selector}`)
+      });
-      );
       const waitFor = page.waitFor('getByClass/foo');
 
       // Set the page content after the waitFor has been started.
@@ -405,5 +403,44 @@ describe('ElementHandle specs', function () {
 
       expect(element).toBeDefined();
     });
+    it('should work when both queryOne and queryAll are registered', async () => {
+      const { page, puppeteer } = getTestState();
+      await page.setContent(
+        '<div id="not-foo"></div><div class="foo"><div id="nested-foo" class="foo"/></div><div class="foo baz">Foo2</div>'
+      );
+      puppeteer.__experimental_registerCustomQueryHandler('getByClass', {
+        queryOne: (element, selector) => element.querySelector(`.${selector}`),
+        queryAll: (element, selector) =>
+          element.querySelectorAll(`.${selector}`),
+      });
+
+      const element = await page.$('getByClass/foo');
+      expect(element).toBeDefined();
+
+      const elements = await page.$$('getByClass/foo');
+      expect(elements.length).toBe(3);
+    });
+    it('should eval when both queryOne and queryAll are registered', async () => {
+      const { page, puppeteer } = getTestState();
+      await page.setContent(
+        '<div id="not-foo"></div><div class="foo">text</div><div class="foo baz">content</div>'
+      );
+      puppeteer.__experimental_registerCustomQueryHandler('getByClass', {
+        queryOne: (element, selector) => element.querySelector(`.${selector}`),
+        queryAll: (element, selector) =>
+          element.querySelectorAll(`.${selector}`),
+      });
+
+      const txtContent = await page.$eval(
+        'getByClass/foo',
+        (div) => div.textContent
+      );
+      expect(txtContent).toBe('text');
+
+      const txtContents = await page.$$eval('getByClass/foo', (divs) =>
+        divs.map((d) => d.textContent).join('')
+      );
+      expect(txtContents).toBe('textcontent');
+    });
   });
 });