diff --git a/src/common/DOMWorld.ts b/src/common/DOMWorld.ts index b15b41cc..bb859dc0 100644 --- a/src/common/DOMWorld.ts +++ b/src/common/DOMWorld.ts @@ -465,6 +465,20 @@ export class DOMWorld { async waitForSelector( selector: string, options: WaitForSelectorOptions + ): Promise { + const { updatedSelector, queryHandler } = getQueryHandlerAndSelector( + selector + ); + return queryHandler.waitFor(this, updatedSelector, options); + } + + /** + * @internal + */ + async waitForSelectorInPage( + queryOne: Function, + selector: string, + options: WaitForSelectorOptions ): Promise { 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 ); diff --git a/src/common/JSHandle.ts b/src/common/JSHandle.ts index bff6b72e..afb6fc7e 100644 --- a/src/common/JSHandle.ts +++ b/src/common/JSHandle.ts @@ -777,15 +777,7 @@ export class ElementHandle< const { updatedSelector, queryHandler } = getQueryHandlerAndSelector( selector ); - - const handle = await this.evaluateHandle( - queryHandler.queryOne, - updatedSelector - ); - const element = handle.asElement(); - if (element) return element; - await handle.dispose(); - return null; + return queryHandler.queryOne(this, updatedSelector); } /** @@ -796,19 +788,7 @@ export class ElementHandle< const { updatedSelector, queryHandler } = getQueryHandlerAndSelector( selector ); - - 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; + return queryHandler.queryAll(this, updatedSelector); } /** @@ -892,15 +872,7 @@ export class ElementHandle< const { updatedSelector, queryHandler } = getQueryHandlerAndSelector( selector ); - const queryHandlerToArray = Function( - 'element', - 'selector', - `return Array.from((${queryHandler.queryAll})(element, selector));` - ) as (...args: unknown[]) => unknown; - const arrayHandle = await this.evaluateHandle( - queryHandlerToArray, - updatedSelector - ); + const arrayHandle = await queryHandler.queryAllArray(this, updatedSelector); const result = await arrayHandle.evaluate< ( elements: Element[], diff --git a/src/common/Puppeteer.ts b/src/common/Puppeteer.ts index 72f538d7..621da963 100644 --- a/src/common/Puppeteer.ts +++ b/src/common/Puppeteer.ts @@ -31,9 +31,9 @@ import { Browser } from './Browser.js'; import { registerCustomQueryHandler, unregisterCustomQueryHandler, - customQueryHandlers, - clearQueryHandlers, - QueryHandler, + customQueryHandlerNames, + clearCustomQueryHandlers, + 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 { - return customQueryHandlers(); + __experimental_customQueryHandlerNames(): string[] { + return customQueryHandlerNames(); } /** @@ -315,6 +315,6 @@ export class Puppeteer { */ // eslint-disable-next-line @typescript-eslint/camelcase __experimental_clearQueryHandlers(): void { - clearQueryHandlers(); + clearCustomQueryHandlers(); } } diff --git a/src/common/QueryHandler.ts b/src/common/QueryHandler.ts index 02814c7f..c1e7551c 100644 --- a/src/common/QueryHandler.ts +++ b/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; + waitFor?: ( + domWorld: DOMWorld, + selector: string, + options: WaitForSelectorOptions + ) => Promise; + queryAll?: ( + element: ElementHandle, + selector: string + ) => Promise; + queryAllArray?: ( + element: ElementHandle, + selector: string + ) => Promise; +} + +export interface CustomQueryHandler { queryOne?: (element: Element | Document, selector: string) => Element | null; queryAll?: ( element: Element | Document, @@ -22,54 +48,103 @@ export interface QueryHandler { ) => Element[] | NodeListOf; } -const _customQueryHandlers = new Map(); +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) => Array.from(res) + ); + return arrayHandle; + }; + } + + return internalHandler; +} + +const _queryHandlers = new Map(); +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 { - return _customQueryHandlers; +export function customQueryHandlerNames(): string[] { + return [..._queryHandlers.keys()]; } -export function clearQueryHandlers(): void { - _customQueryHandlers.clear(); +export function clearCustomQueryHandlers(): void { + _queryHandlers.clear(); } export function getQueryHandlerAndSelector( selector: string -): { updatedSelector: string; queryHandler: QueryHandler } { - const defaultHandler = { - queryOne: (element: Element, selector: string) => - element.querySelector(selector), - queryAll: (element: Element, selector: string) => - element.querySelectorAll(selector), - }; +): { updatedSelector: string; queryHandler: InternalQueryHandler } { 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` diff --git a/test/queryselector.spec.ts b/test/queryselector.spec.ts index 854fe851..2b36c0cb 100644 --- a/test/queryselector.spec.ts +++ b/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. 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( + '
hello
beautiful
world!
' + ); + 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( + '
2
2
1
3
' + ); + 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( + '
A

B
' + ); + 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( + 'A
B' + ); + 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( + '
hello
beautiful
world!
' + ); + 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( + '
hello
beautiful
world!
' + ); + 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( + '
2
2
1
3
' + ); + 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); + }); + }); });