From 4a47867a2444e28257dec5113b647fbe3cbbf54e Mon Sep 17 00:00:00 2001 From: Paul Lewis Date: Thu, 30 Apr 2020 12:45:52 +0100 Subject: [PATCH] feat: add support for string-based custom queries (#5753) --- src/DOMWorld.ts | 29 ++++++++++++--- src/JSHandle.ts | 24 ++++++------- src/Puppeteer.js | 24 +++++++++++++ src/QueryHandler.ts | 74 ++++++++++++++++++++++++++++++++++++++ test/elementhandle.spec.js | 73 +++++++++++++++++++++++++++++++++++++ 5 files changed, 207 insertions(+), 17 deletions(-) create mode 100644 src/QueryHandler.ts diff --git a/src/DOMWorld.ts b/src/DOMWorld.ts index 92143506a4bd7..bc184610c7053 100644 --- a/src/DOMWorld.ts +++ b/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 { @@ -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); diff --git a/src/JSHandle.ts b/src/JSHandle.ts index c78eec3bd14e8..61f53db40eccc 100644 --- a/src/JSHandle.ts +++ b/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 { - 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>} */ async $$(selector: string): Promise { - 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(selector: string, pageFunction: Function | string, ...args: unknown[]): Promise { - 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(pageFunction, ...args); await arrayHandle.dispose(); return result; diff --git a/src/Puppeteer.js b/src/Puppeteer.js index 8f84cacfc5999..023fc24fd5fda 100644 --- a/src/Puppeteer.js +++ b/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(); + } }; diff --git a/src/QueryHandler.ts b/src/QueryHandler.ts new file mode 100644 index 0000000000000..98f6b1e896c35 --- /dev/null +++ b/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; +} + +const _customQueryHandlers = new Map(); + +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 { + 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 +}; diff --git a/test/elementhandle.spec.js b/test/elementhandle.spec.js index 4ec76ca0004a8..7b53f30a5d0f8 100644 --- a/test/elementhandle.spec.js +++ b/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('
'); + + // 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('
Foo1
Foo2
'); + 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('
Foo1
Foo2
'); + 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('
Foo1
'); + 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('
Foo1
'); + const element = await waitFor; + + expect(element).toBeDefined(); + }); + }); });