feat: add support for string-based custom queries (#5753)

src/DOMWorld.ts

@@ -23,6 +23,11 @@ import {ExecutionContext} from './ExecutionContext';
 import {TimeoutSettings} from './TimeoutSettings';
 import {MouseButtonInput} from './Input';
 import {FrameManager, Frame} from './FrameManager';
+import {getQueryHandlerAndSelector, QueryHandler} from './QueryHandler';
+
+// This predicateQueryHandler is declared here so that TypeScript knows about it
+// when it is used in the predicate function below.
+declare const predicateQueryHandler: QueryHandler;
 
 const readFileAsync = helper.promisify(fs.readFile);
 
@@ -364,7 +369,7 @@ export class DOMWorld {
       polling = 'raf',
       timeout = this._timeoutSettings.timeout(),
     } = options;
-    return new WaitTask(this, pageFunction, 'function', polling, timeout, ...args).promise;
+    return new WaitTask(this, pageFunction, undefined, 'function', polling, timeout, ...args).promise;
   }
 
   async title(): Promise<string> {
@@ -379,7 +384,8 @@ export class DOMWorld {
     } = options;
     const polling = waitForVisible || waitForHidden ? 'raf' : 'mutation';
     const title = `${isXPath ? 'XPath' : 'selector'} "${selectorOrXPath}"${waitForHidden ? ' to be hidden' : ''}`;
-    const waitTask = new WaitTask(this, predicate, title, polling, timeout, selectorOrXPath, isXPath, waitForVisible, waitForHidden);
+    const {updatedSelector, queryHandler} = getQueryHandlerAndSelector(selectorOrXPath, (element, selector) => document.querySelector(selector));
+    const waitTask = new WaitTask(this, predicate, queryHandler, title, polling, timeout, updatedSelector, isXPath, waitForVisible, waitForHidden);
     const handle = await waitTask.promise;
     if (!handle.asElement()) {
       await handle.dispose();
@@ -397,7 +403,7 @@ export class DOMWorld {
     function predicate(selectorOrXPath: string, isXPath: boolean, waitForVisible: boolean, waitForHidden: boolean): Node | null | boolean {
       const node = isXPath
         ? document.evaluate(selectorOrXPath, document, null, XPathResult.FIRST_ORDERED_NODE_TYPE, null).singleNodeValue
-        : document.querySelector(selectorOrXPath);
+        : predicateQueryHandler ? predicateQueryHandler(document, selectorOrXPath) as Element : document.querySelector(selectorOrXPath);
       if (!node)
         return waitForHidden;
       if (!waitForVisible && !waitForHidden)
@@ -430,7 +436,7 @@ class WaitTask {
   _timeoutTimer?: NodeJS.Timeout;
   _terminated = false;
 
-  constructor(domWorld: DOMWorld, predicateBody: Function | string, title: string, polling: string | number, timeout: number, ...args: unknown[]) {
+  constructor(domWorld: DOMWorld, predicateBody: Function | string, predicateQueryHandlerBody: Function | string | undefined, title: string, polling: string | number, timeout: number, ...args: unknown[]) {
     if (helper.isString(polling))
       assert(polling === 'raf' || polling === 'mutation', 'Unknown polling option: ' + polling);
     else if (helper.isNumber(polling))
@@ -438,10 +444,23 @@ class WaitTask {
     else
       throw new Error('Unknown polling options: ' + polling);
 
+    function getPredicateBody(predicateBody: Function | string, predicateQueryHandlerBody: Function | string) {
+      if (helper.isString(predicateBody))
+        return `return (${predicateBody});`;
+      if (predicateQueryHandlerBody) {
+        return `
+          return (function wrapper(args) {
+            const predicateQueryHandler = ${predicateQueryHandlerBody};
+            return (${predicateBody})(...args);
+          })(args);`;
+      }
+      return `return (${predicateBody})(...args);`;
+    }
+
     this._domWorld = domWorld;
     this._polling = polling;
     this._timeout = timeout;
-    this._predicateBody = helper.isString(predicateBody) ? 'return (' + predicateBody + ')' : 'return (' + predicateBody + ')(...args)';
+    this._predicateBody = getPredicateBody(predicateBody, predicateQueryHandlerBody);
     this._args = args;
     this._runCount = 0;
     domWorld._waitTasks.add(this);

src/JSHandle.ts

@@ -19,6 +19,7 @@ import {ExecutionContext} from './ExecutionContext';
 import {CDPSession} from './Connection';
 import {KeyInput} from './USKeyboardLayout';
 import {FrameManager, Frame} from './FrameManager';
+import {getQueryHandlerAndSelector} from './QueryHandler';
 
 interface BoxModel {
   content: Array<{x: number; y: number}>;
@@ -427,10 +428,10 @@ export class ElementHandle extends JSHandle {
   }
 
   async $(selector: string): Promise<ElementHandle | null> {
-    const handle = await this.evaluateHandle(
-        (element, selector) => element.querySelector(selector),
-        selector
-    );
+    const defaultHandler = (element: Element, selector: string) => element.querySelector(selector);
+    const {updatedSelector, queryHandler} = getQueryHandlerAndSelector(selector, defaultHandler);
+
+    const handle = await this.evaluateHandle(queryHandler, updatedSelector);
     const element = handle.asElement();
     if (element)
       return element;
@@ -443,10 +444,10 @@ export class ElementHandle extends JSHandle {
    * @return {!Promise<!Array<!ElementHandle>>}
    */
   async $$(selector: string): Promise<ElementHandle[]> {
-    const arrayHandle = await this.evaluateHandle(
-        (element, selector) => element.querySelectorAll(selector),
-        selector
-    );
+    const defaultHandler = (element: Element, selector: string) => element.querySelectorAll(selector);
+    const {updatedSelector, queryHandler} = getQueryHandlerAndSelector(selector, defaultHandler);
+
+    const arrayHandle = await this.evaluateHandle(queryHandler, updatedSelector);
     const properties = await arrayHandle.getProperties();
     await arrayHandle.dispose();
     const result = [];
@@ -468,11 +469,10 @@ export class ElementHandle extends JSHandle {
   }
 
   async $$eval<ReturnType extends any>(selector: string, pageFunction: Function | string, ...args: unknown[]): Promise<ReturnType> {
-    const arrayHandle = await this.evaluateHandle(
-        (element, selector) => Array.from(element.querySelectorAll(selector)),
-        selector
-    );
+    const defaultHandler = (element: Element, selector: string) => Array.from(element.querySelectorAll(selector));
+    const {updatedSelector, queryHandler} = getQueryHandlerAndSelector(selector, defaultHandler);
 
+    const arrayHandle = await this.evaluateHandle(queryHandler, updatedSelector);
     const result = await arrayHandle.evaluate<ReturnType>(pageFunction, ...args);
     await arrayHandle.dispose();
     return result;

src/Puppeteer.js

@@ -20,6 +20,7 @@ const DeviceDescriptors = require('./DeviceDescriptors');
 // Import used as typedef
 // eslint-disable-next-line no-unused-vars
 const {Browser} = require('./Browser');
+const QueryHandler = require('./QueryHandler');
 
 module.exports = class {
   /**
@@ -147,4 +148,27 @@ module.exports = class {
   createBrowserFetcher(options) {
     return new BrowserFetcher(this._projectRoot, options);
   }
+
+  /**
+   * @param {string} name
+   * @param {!Function} queryHandler
+   */
+  __experimental_registerCustomQueryHandler(name, queryHandler) {
+    QueryHandler.registerCustomQueryHandler(name, queryHandler);
+  }
+
+  /**
+   * @param {string} name
+   */
+  __experimental_unregisterCustomQueryHandler(name) {
+    QueryHandler.unregisterCustomQueryHandler(name);
+  }
+
+  __experimental_customQueryHandlers() {
+    return QueryHandler.customQueryHandlers();
+  }
+
+  __experimental_clearQueryHandlers() {
+    QueryHandler.clearQueryHandlers();
+  }
 };

src/QueryHandler.ts

@@ -0,0 +1,74 @@
+/**
+ * Copyright 2020 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.
+ */
+
+export interface QueryHandler {
+  (element: Element | Document, selector: string): Element | Element[] | NodeListOf<Element>;
+}
+
+const _customQueryHandlers = new Map<string, QueryHandler>();
+
+export function registerCustomQueryHandler(name: string, handler: Function): void {
+  if (_customQueryHandlers.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 as QueryHandler);
+}
+
+/**
+ * @param {string} name
+ */
+export function unregisterCustomQueryHandler(name: string): void {
+  _customQueryHandlers.delete(name);
+}
+
+export function customQueryHandlers(): Map<string, QueryHandler> {
+  return _customQueryHandlers;
+}
+
+export function clearQueryHandlers(): void {
+  _customQueryHandlers.clear();
+}
+
+export function getQueryHandlerAndSelector(selector: string, defaultQueryHandler: QueryHandler):
+    { updatedSelector: string; queryHandler: QueryHandler} {
+  const hasCustomQueryHandler = /^[a-zA-Z]+\//.test(selector);
+  if (!hasCustomQueryHandler)
+    return {updatedSelector: selector, queryHandler: defaultQueryHandler};
+
+  const index = selector.indexOf('/');
+  const name = selector.slice(0, index);
+  const updatedSelector = selector.slice(index + 1);
+  const queryHandler = customQueryHandlers().get(name);
+  if (!queryHandler)
+    throw new Error(`Query set to use "${name}", but no query handler of that name was found`);
+
+  return {
+    updatedSelector,
+    queryHandler
+  };
+}
+
+module.exports = {
+  registerCustomQueryHandler,
+  unregisterCustomQueryHandler,
+  customQueryHandlers,
+  getQueryHandlerAndSelector,
+  clearQueryHandlers
+};

test/elementhandle.spec.js

@@ -248,4 +248,77 @@ describe('ElementHandle specs', function() {
       }
     });
   });
+
+  describe('Custom queries', function() {
+    this.afterEach(() => {
+      const {puppeteer} = getTestState();
+      puppeteer.__experimental_clearQueryHandlers();
+    });
+    it('should register and unregister', async() => {
+      const {page, puppeteer} = getTestState();
+      await page.setContent('<div id="not-foo"></div><div id="foo"></div>');
+
+      // Register.
+      puppeteer.__experimental_registerCustomQueryHandler('getById', (element, selector) => document.querySelector(`[id="${selector}"]`));
+      const element = await page.$('getById/foo');
+      expect(await page.evaluate(element => element.id, element)).toBe('foo');
+
+      // Unregister.
+      puppeteer.__experimental_unregisterCustomQueryHandler('getById');
+      try {
+        await page.$('getById/foo');
+        expect.fail('Custom query handler not set - throw expected');
+      } catch (error) {
+        expect(error).toStrictEqual(new Error('Query set to use "getById", but no query handler of that name was found'));
+      }
+    });
+    it('should throw with invalid query names', () => {
+      try {
+        const {puppeteer} = getTestState();
+        puppeteer.__experimental_registerCustomQueryHandler('1/2/3', (element, selector) => {});
+        expect.fail('Custom query handler name was invalid - throw expected');
+      } catch (error) {
+        expect(error).toStrictEqual(new Error('Custom query handler names may only contain [a-zA-Z]'));
+      }
+    });
+    it('should work for multiple elements', async() => {
+      const {page, puppeteer} = getTestState();
+      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}`));
+      const elements = await page.$$('getByClass/foo');
+      const classNames = await Promise.all(elements.map(async element => await page.evaluate(element => element.className, element)));
+
+      expect(classNames).toStrictEqual(['foo', 'foo baz']);
+    });
+    it('should eval correctly', async() => {
+      const {page, puppeteer} = getTestState();
+      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}`));
+      const elements = await page.$$eval('getByClass/foo', divs => divs.length);
+
+      expect(elements).toBe(2);
+    });
+    it('should wait correctly with waitForSelector', async() => {
+      const {page, puppeteer} = getTestState();
+      puppeteer.__experimental_registerCustomQueryHandler('getByClass', (element, selector) => element.querySelector(`.${selector}`));
+      const waitFor = page.waitForSelector('getByClass/foo');
+
+      // Set the page content after the waitFor has been started.
+      await page.setContent('<div id="not-foo"></div><div class="foo">Foo1</div>');
+      const element = await waitFor;
+
+      expect(element).toBeDefined();
+    });
+    it('should wait correctly with waitFor', async() => {
+      const {page, puppeteer} = getTestState();
+      puppeteer.__experimental_registerCustomQueryHandler('getByClass', (element, selector) => element.querySelector(`.${selector}`));
+      const waitFor = page.waitFor('getByClass/foo');
+
+      // Set the page content after the waitFor has been started.
+      await page.setContent('<div id="not-foo"></div><div class="foo">Foo1</div>');
+      const element = await waitFor;
+
+      expect(element).toBeDefined();
+    });
+  });
 });