diff --git a/src/common/DOMWorld.ts b/src/common/DOMWorld.ts index 1a64a7edca606..9e8d65f2bdaa3 100644 --- a/src/common/DOMWorld.ts +++ b/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; /** * @public @@ -506,16 +509,13 @@ export class DOMWorld { const title = `${isXPath ? 'XPath' : 'selector'} "${selectorOrXPath}"${ waitForHidden ? ' to be hidden' : '' }`; - const { - updatedSelector, - queryHandler, - } = getQueryHandlerAndSelector(selectorOrXPath, (element, selector) => - document.querySelector(selector) + const { updatedSelector, queryHandler } = getQueryHandlerAndSelector( + selectorOrXPath ); const waitTask = new WaitTask( this, predicate, - queryHandler, + queryHandler.queryOne, title, polling, timeout, diff --git a/src/common/JSHandle.ts b/src/common/JSHandle.ts index 806ad81429729..bff6b72e6adff 100644 --- a/src/common/JSHandle.ts +++ b/src/common/JSHandle.ts @@ -774,14 +774,14 @@ export class ElementHandle< * the return value resolves to `null`. */ async $(selector: string): Promise { - const defaultHandler = (element: Element, selector: string) => - element.querySelector(selector); const { updatedSelector, queryHandler } = getQueryHandlerAndSelector( - selector, - defaultHandler + selector ); - 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 { - const defaultHandler = (element: Element, selector: string) => - element.querySelectorAll(selector); const { updatedSelector, queryHandler } = getQueryHandlerAndSelector( - selector, - defaultHandler + selector ); - const arrayHandle = await this.evaluateHandle( - queryHandler, + const handles = await this.evaluateHandle( + queryHandler.queryAll, updatedSelector ); - const properties = await arrayHandle.getProperties(); - await arrayHandle.dispose(); + const properties = await handles.getProperties(); + 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 - * `elementHandle.evaluate`. If evalute returns an element it will return an + * This `as` is a little unfortunate but helps TS understand the behavior of + * `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 if it is an ElementHandle, or leave it alone as @@ -892,15 +889,16 @@ export class ElementHandle< ) => ReturnType | Promise, ...args: SerializableOrJSHandle[] ): Promise> { - const defaultHandler = (element: Element, selector: string) => - Array.from(element.querySelectorAll(selector)); const { updatedSelector, queryHandler } = getQueryHandlerAndSelector( - selector, - defaultHandler + selector ); - + 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 >(pageFunction, ...args); await arrayHandle.dispose(); - /* This as exists for the same reason as the `as` in $eval above. - * See the comment there for a ful explanation. + /* This `as` exists for the same reason as the `as` in $eval above. + * See the comment there for a full explanation. */ return result as WrapElementHandle; } diff --git a/src/common/QueryHandler.ts b/src/common/QueryHandler.ts index 7f86c2425fa3a..02814c7f969f3 100644 --- a/src/common/QueryHandler.ts +++ b/src/common/QueryHandler.ts @@ -15,17 +15,18 @@ */ export interface QueryHandler { - (element: Element | Document, selector: string): - | Element - | Element[] - | NodeListOf; + queryOne?: (element: Element | Document, selector: string) => Element | null; + queryAll?: ( + element: Element | Document, + selector: string + ) => Element[] | NodeListOf; } const _customQueryHandlers = new Map(); 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, - defaultQueryHandler: QueryHandler + selector: string ): { 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); diff --git a/test/elementhandle.spec.ts b/test/elementhandle.spec.ts index 5bb74a3dbf397..79d3225196da5 100644 --- a/test/elementhandle.spec.ts +++ b/test/elementhandle.spec.ts @@ -293,10 +293,10 @@ describe('ElementHandle specs', function () { await page.setContent('
'); // Register. - puppeteer.__experimental_registerCustomQueryHandler( - 'getById', - (element, selector) => document.querySelector(`[id="${selector}"]`) - ); + puppeteer.__experimental_registerCustomQueryHandler('getById', { + queryOne: (element, 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( '
Foo1
Foo2
' ); - puppeteer.__experimental_registerCustomQueryHandler( - 'getByClass', - (element, selector) => document.querySelectorAll(`.${selector}`) - ); + puppeteer.__experimental_registerCustomQueryHandler('getByClass', { + queryAll: (element, 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( '
Foo1
Foo2
' ); - puppeteer.__experimental_registerCustomQueryHandler( - 'getByClass', - (element, selector) => document.querySelectorAll(`.${selector}`) - ); + puppeteer.__experimental_registerCustomQueryHandler('getByClass', { + queryAll: (element, 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( - 'getByClass', - (element, selector) => element.querySelector(`.${selector}`) - ); + puppeteer.__experimental_registerCustomQueryHandler('getByClass', { + queryOne: (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( - 'getByClass', - (element, selector) => element.querySelector(`.${selector}`) - ); + puppeteer.__experimental_registerCustomQueryHandler('getByClass', { + queryOne: (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( + '
Foo2
' + ); + 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( + '
text
content
' + ); + 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'); + }); }); });