feat(a11y-query): introduce internal handlers (#6437)

This commit changes the internal representation of query handlers to contain Puppeteer-level code instead of page functions.
The interface `CustomQueryHandler` is introduced for user-defined query handlers. When a `CustomQueryHandler` is registered using  `registerCustomQueryHandler` a corresponding Puppeteer-level handler is created through `makeQueryHandler` by wrapping the page functions as appropriate.
The internal query handlers (defined by the interface `QueryHandler`) contain two new functions: `waitFor` and `queryAllArray`.
- `waitFor` allows page-based handlers to make use of the `WaitTask`-backed implementation in `DOMWorld`, whereas purely Puppeteer-based handlers can define an alternative approach instead.
- `queryAllArray` is similar to `queryAll` but with a slightly different interface; it returns a `JSHandle` to an array with the results as opposed to an array of `ElementHandle`. It is used by `$$eval`. 

After this change, we can introduce built-in query handlers that are not executed in the page context (#6307).

src/common/DOMWorld.ts

@@ -465,6 +465,20 @@ export class DOMWorld {
   async waitForSelector(
     selector: string,
     options: WaitForSelectorOptions
+  ): Promise<ElementHandle | null> {
+    const { updatedSelector, queryHandler } = getQueryHandlerAndSelector(
+      selector
+    );
+    return queryHandler.waitFor(this, updatedSelector, options);
+  }
+
+  /**
+   * @internal
+   */
+  async waitForSelectorInPage(
+    queryOne: Function,
+    selector: string,
+    options: WaitForSelectorOptions
   ): Promise<ElementHandle | null> {
     const {
       visible: waitForVisible = false,
@@ -475,9 +489,6 @@ export class DOMWorld {
     const title = `selector \`${selector}\`${
       waitForHidden ? ' to be hidden' : ''
     }`;
-    const { updatedSelector, queryHandler } = getQueryHandlerAndSelector(
-      selector
-    );
     function predicate(
       selector: string,
       waitForVisible: boolean,
@@ -490,11 +501,11 @@ export class DOMWorld {
     }
     const waitTask = new WaitTask(
       this,
-      this._makePredicateString(predicate, queryHandler.queryOne),
+      this._makePredicateString(predicate, queryOne),
       title,
       polling,
       timeout,
-      updatedSelector,
+      selector,
       waitForVisible,
       waitForHidden
     );

src/common/JSHandle.ts

@@ -777,15 +777,7 @@ export class ElementHandle<
     const { updatedSelector, queryHandler } = getQueryHandlerAndSelector(
       selector
     );
-
+    return queryHandler.queryOne(this, updatedSelector);
-    const handle = await this.evaluateHandle(
-      queryHandler.queryOne,
-      updatedSelector
-    );
-    const element = handle.asElement();
-    if (element) return element;
-    await handle.dispose();
-    return null;
   }
 
   /**
@@ -796,19 +788,7 @@ export class ElementHandle<
     const { updatedSelector, queryHandler } = getQueryHandlerAndSelector(
       selector
     );
-
+    return queryHandler.queryAll(this, updatedSelector);
-    const handles = await this.evaluateHandle(
-      queryHandler.queryAll,
-      updatedSelector
-    );
-    const properties = await handles.getProperties();
-    await handles.dispose();
-    const result = [];
-    for (const property of properties.values()) {
-      const elementHandle = property.asElement();
-      if (elementHandle) result.push(elementHandle);
-    }
-    return result;
   }
 
   /**
@@ -892,15 +872,7 @@ export class ElementHandle<
     const { updatedSelector, queryHandler } = getQueryHandlerAndSelector(
       selector
     );
-    const queryHandlerToArray = Function(
+    const arrayHandle = await queryHandler.queryAllArray(this, updatedSelector);
-      'element',
-      'selector',
-      `return Array.from((${queryHandler.queryAll})(element, selector));`
-    ) as (...args: unknown[]) => unknown;
-    const arrayHandle = await this.evaluateHandle(
-      queryHandlerToArray,
-      updatedSelector
-    );
     const result = await arrayHandle.evaluate<
       (
         elements: Element[],

src/common/Puppeteer.ts

@@ -31,9 +31,9 @@ import { Browser } from './Browser.js';
 import {
   registerCustomQueryHandler,
   unregisterCustomQueryHandler,
-  customQueryHandlers,
+  customQueryHandlerNames,
-  clearQueryHandlers,
+  clearCustomQueryHandlers,
-  QueryHandler,
+  CustomQueryHandler,
 } from './QueryHandler.js';
 import { PUPPETEER_REVISIONS } from '../revisions.js';
 
@@ -289,7 +289,7 @@ export class Puppeteer {
   // eslint-disable-next-line @typescript-eslint/camelcase
   __experimental_registerCustomQueryHandler(
     name: string,
-    queryHandler: QueryHandler
+    queryHandler: CustomQueryHandler
   ): void {
     registerCustomQueryHandler(name, queryHandler);
   }
@@ -306,8 +306,8 @@ export class Puppeteer {
    * @internal
    */
   // eslint-disable-next-line @typescript-eslint/camelcase
-  __experimental_customQueryHandlers(): Map<string, QueryHandler> {
+  __experimental_customQueryHandlerNames(): string[] {
-    return customQueryHandlers();
+    return customQueryHandlerNames();
   }
 
   /**
@@ -315,6 +315,6 @@ export class Puppeteer {
    */
   // eslint-disable-next-line @typescript-eslint/camelcase
   __experimental_clearQueryHandlers(): void {
-    clearQueryHandlers();
+    clearCustomQueryHandlers();
   }
 }

src/common/QueryHandler.ts

@@ -14,7 +14,33 @@
  * limitations under the License.
  */
 
-export interface QueryHandler {
+import { WaitForSelectorOptions, DOMWorld } from './DOMWorld.js';
+import { ElementHandle, JSHandle } from './JSHandle.js';
+
+/**
+ * @internal
+ */
+interface InternalQueryHandler {
+  queryOne?: (
+    element: ElementHandle,
+    selector: string
+  ) => Promise<ElementHandle | null>;
+  waitFor?: (
+    domWorld: DOMWorld,
+    selector: string,
+    options: WaitForSelectorOptions
+  ) => Promise<ElementHandle | null>;
+  queryAll?: (
+    element: ElementHandle,
+    selector: string
+  ) => Promise<ElementHandle[]>;
+  queryAllArray?: (
+    element: ElementHandle,
+    selector: string
+  ) => Promise<JSHandle>;
+}
+
+export interface CustomQueryHandler {
   queryOne?: (element: Element | Document, selector: string) => Element | null;
   queryAll?: (
     element: Element | Document,
@@ -22,54 +48,103 @@ export interface QueryHandler {
   ) => Element[] | NodeListOf<Element>;
 }
 
-const _customQueryHandlers = new Map<string, QueryHandler>();
+function makeQueryHandler(handler: CustomQueryHandler): InternalQueryHandler {
+  const internalHandler: InternalQueryHandler = {};
+
+  if (handler.queryOne) {
+    internalHandler.queryOne = async (element, selector) => {
+      const jsHandle = await element.evaluateHandle(handler.queryOne, selector);
+      const elementHandle = jsHandle.asElement();
+      if (elementHandle) return elementHandle;
+      await jsHandle.dispose();
+      return null;
+    };
+    internalHandler.waitFor = (
+      domWorld: DOMWorld,
+      selector: string,
+      options: WaitForSelectorOptions
+    ) => domWorld.waitForSelectorInPage(handler.queryOne, selector, options);
+  }
+
+  if (handler.queryAll) {
+    internalHandler.queryAll = async (element, selector) => {
+      const jsHandle = await element.evaluateHandle(handler.queryAll, selector);
+      const properties = await jsHandle.getProperties();
+      await jsHandle.dispose();
+      const result = [];
+      for (const property of properties.values()) {
+        const elementHandle = property.asElement();
+        if (elementHandle) result.push(elementHandle);
+      }
+      return result;
+    };
+    internalHandler.queryAllArray = async (element, selector) => {
+      const resultHandle = await element.evaluateHandle(
+        handler.queryAll,
+        selector
+      );
+      const arrayHandle = await resultHandle.evaluateHandle(
+        (res: Element[] | NodeListOf<Element>) => Array.from(res)
+      );
+      return arrayHandle;
+    };
+  }
+
+  return internalHandler;
+}
+
+const _queryHandlers = new Map<string, InternalQueryHandler>();
+const _defaultHandler = makeQueryHandler({
+  queryOne: (element: Element, selector: string) =>
+    element.querySelector(selector),
+  queryAll: (element: Element, selector: string) =>
+    element.querySelectorAll(selector),
+});
 
 export function registerCustomQueryHandler(
   name: string,
-  handler: QueryHandler
+  handler: CustomQueryHandler
 ): void {
-  if (_customQueryHandlers.get(name))
+  if (_queryHandlers.get(name))
     throw new Error(`A custom 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, handler);
+  const internalHandler = makeQueryHandler(handler);
+
+  _queryHandlers.set(name, internalHandler);
 }
 
 /**
  * @param {string} name
  */
 export function unregisterCustomQueryHandler(name: string): void {
-  _customQueryHandlers.delete(name);
+  if (_queryHandlers.has(name)) {
+    _queryHandlers.delete(name);
+  }
 }
 
-export function customQueryHandlers(): Map<string, QueryHandler> {
+export function customQueryHandlerNames(): string[] {
-  return _customQueryHandlers;
+  return [..._queryHandlers.keys()];
 }
 
-export function clearQueryHandlers(): void {
+export function clearCustomQueryHandlers(): void {
-  _customQueryHandlers.clear();
+  _queryHandlers.clear();
 }
 
 export function getQueryHandlerAndSelector(
   selector: string
-): { updatedSelector: string; queryHandler: QueryHandler } {
+): { updatedSelector: string; queryHandler: InternalQueryHandler } {
-  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: defaultHandler };
+    return { updatedSelector: selector, queryHandler: _defaultHandler };
 
   const index = selector.indexOf('/');
   const name = selector.slice(0, index);
   const updatedSelector = selector.slice(index + 1);
-  const queryHandler = customQueryHandlers().get(name);
+  const queryHandler = _queryHandlers.get(name);
   if (!queryHandler)
     throw new Error(
       `Query set to use "${name}", but no query handler of that name was found`

test/queryselector.spec.ts

@@ -19,6 +19,7 @@ import {
   setupTestBrowserHooks,
   setupTestPageAndContextHooks,
 } from './mocha-utils'; // eslint-disable-line import/extensions
+import { CustomQueryHandler } from '../lib/cjs/puppeteer/common/QueryHandler.js';
 
 describe('querySelector', function () {
   setupTestBrowserHooks();
@@ -67,6 +68,9 @@ describe('querySelector', function () {
     });
   });
 
+  // The tests for $$eval are repeated later in this file in the test group 'QueryAll'.
+  // This is done to also test a query handler where QueryAll returns an Element[]
+  // as opposed to NodeListOf<Element>.
   describe('Page.$$eval', function () {
     it('should work', async () => {
       const { page } = getTestState();
@@ -77,6 +81,52 @@ describe('querySelector', function () {
       const divsCount = await page.$$eval('div', (divs) => divs.length);
       expect(divsCount).toBe(3);
     });
+    it('should accept extra arguments', async () => {
+      const { page } = getTestState();
+      await page.setContent(
+        '<div>hello</div><div>beautiful</div><div>world!</div>'
+      );
+      const divsCountPlus5 = await page.$$eval(
+        'div',
+        (divs, two: number, three: number) => divs.length + two + three,
+        2,
+        3
+      );
+      expect(divsCountPlus5).toBe(8);
+    });
+    it('should accept ElementHandles as arguments', async () => {
+      const { page } = getTestState();
+      await page.setContent(
+        '<section>2</section><section>2</section><section>1</section><div>3</div>'
+      );
+      const divHandle = await page.$('div');
+      const sum = await page.$$eval(
+        'section',
+        (sections, div: HTMLElement) =>
+          sections.reduce(
+            (acc, section) => acc + Number(section.textContent),
+            0
+          ) + Number(div.textContent),
+        divHandle
+      );
+      expect(sum).toBe(8);
+    });
+    it('should handle many elements', async () => {
+      const { page } = getTestState();
+      await page.evaluate(
+        `
+        for (var i = 0; i <= 1000; i++) {
+            const section = document.createElement('section');
+            section.textContent = i;
+            document.body.appendChild(section);
+        }
+        `
+      );
+      const sum = await page.$$eval('section', (sections) =>
+        sections.reduce((acc, section) => acc + Number(section.textContent), 0)
+      );
+      expect(sum).toBe(500500);
+    });
   });
 
   describe('Page.$', function () {
@@ -312,4 +362,100 @@ describe('querySelector', function () {
       expect(second).toEqual([]);
     });
   });
+
+  // This is the same tests for `$$eval` and `$$` as above, but with a queryAll
+  // handler that returns an array instead of a list of nodes.
+  describe('QueryAll', function () {
+    const handler: CustomQueryHandler = {
+      queryAll: (element: Element, selector: string) =>
+        Array.from(element.querySelectorAll(selector)),
+    };
+    before(() => {
+      const { puppeteer } = getTestState();
+      puppeteer.__experimental_registerCustomQueryHandler('allArray', handler);
+    });
+    it('$$ should query existing elements', async () => {
+      const { page } = getTestState();
+
+      await page.setContent(
+        '<html><body><div>A</div><br/><div>B</div></body></html>'
+      );
+      const html = await page.$('html');
+      const elements = await html.$$('allArray/div');
+      expect(elements.length).toBe(2);
+      const promises = elements.map((element) =>
+        page.evaluate((e: HTMLElement) => e.textContent, element)
+      );
+      expect(await Promise.all(promises)).toEqual(['A', 'B']);
+    });
+
+    it('$$ should return empty array for non-existing elements', async () => {
+      const { page } = getTestState();
+
+      await page.setContent(
+        '<html><body><span>A</span><br/><span>B</span></body></html>'
+      );
+      const html = await page.$('html');
+      const elements = await html.$$('allArray/div');
+      expect(elements.length).toBe(0);
+    });
+    it('$$eval should work', async () => {
+      const { page } = getTestState();
+
+      await page.setContent(
+        '<div>hello</div><div>beautiful</div><div>world!</div>'
+      );
+      const divsCount = await page.$$eval(
+        'allArray/div',
+        (divs) => divs.length
+      );
+      expect(divsCount).toBe(3);
+    });
+    it('$$eval should accept extra arguments', async () => {
+      const { page } = getTestState();
+      await page.setContent(
+        '<div>hello</div><div>beautiful</div><div>world!</div>'
+      );
+      const divsCountPlus5 = await page.$$eval(
+        'allArray/div',
+        (divs, two: number, three: number) => divs.length + two + three,
+        2,
+        3
+      );
+      expect(divsCountPlus5).toBe(8);
+    });
+    it('$$eval should accept ElementHandles as arguments', async () => {
+      const { page } = getTestState();
+      await page.setContent(
+        '<section>2</section><section>2</section><section>1</section><div>3</div>'
+      );
+      const divHandle = await page.$('div');
+      const sum = await page.$$eval(
+        'allArray/section',
+        (sections, div: HTMLElement) =>
+          sections.reduce(
+            (acc, section) => acc + Number(section.textContent),
+            0
+          ) + Number(div.textContent),
+        divHandle
+      );
+      expect(sum).toBe(8);
+    });
+    it('$$eval should handle many elements', async () => {
+      const { page } = getTestState();
+      await page.evaluate(
+        `
+        for (var i = 0; i <= 1000; i++) {
+            const section = document.createElement('section');
+            section.textContent = i;
+            document.body.appendChild(section);
+        }
+        `
+      );
+      const sum = await page.$$eval('allArray/section', (sections) =>
+        sections.reduce((acc, section) => acc + Number(section.textContent), 0)
+      );
+      expect(sum).toBe(500500);
+    });
+  });
 });