From e8f25e403ffb62bc4bcd160bc335c98546ff0854 Mon Sep 17 00:00:00 2001 From: jrandolf <101637635+jrandolf@users.noreply.github.com> Date: Wed, 15 Feb 2023 07:33:18 -0800 Subject: [PATCH] refactor: custom query handlers and global bindings (#9678) --- .../src/common/AriaQueryHandler.ts | 31 +-- packages/puppeteer-core/src/common/Binding.ts | 13 +- .../src/common/CustomQueryHandler.ts | 181 +++++++++++++++--- .../src/common/ElementHandle.ts | 12 +- .../src/common/ExecutionContext.ts | 48 +++-- packages/puppeteer-core/src/common/Frame.ts | 4 +- .../src/common/GetQueryHandler.ts | 12 +- .../src/common/IsolatedWorld.ts | 3 +- .../puppeteer-core/src/common/Puppeteer.ts | 25 +-- .../puppeteer-core/src/common/QueryHandler.ts | 72 +++---- .../src/common/ScriptInjector.ts | 49 +++++ .../puppeteer-core/src/common/WaitTask.ts | 3 +- .../puppeteer-core/src/common/bidi/Page.ts | 3 +- packages/puppeteer-core/src/common/util.ts | 26 --- .../src/injected/ARIAQuerySelector.ts | 31 +++ .../src/injected/CustomQuerySelector.ts | 68 +++++++ .../puppeteer-core/src/injected/injected.ts | 20 +- packages/puppeteer-core/src/util/Function.ts | 63 ++++++ test/src/elementhandle.spec.ts | 21 ++ 19 files changed, 513 insertions(+), 172 deletions(-) create mode 100644 packages/puppeteer-core/src/common/ScriptInjector.ts create mode 100644 packages/puppeteer-core/src/injected/ARIAQuerySelector.ts create mode 100644 packages/puppeteer-core/src/injected/CustomQuerySelector.ts diff --git a/packages/puppeteer-core/src/common/AriaQueryHandler.ts b/packages/puppeteer-core/src/common/AriaQueryHandler.ts index 066c1a619c7..f16b6de9a38 100644 --- a/packages/puppeteer-core/src/common/AriaQueryHandler.ts +++ b/packages/puppeteer-core/src/common/AriaQueryHandler.ts @@ -19,8 +19,6 @@ import {Protocol} from 'devtools-protocol'; import {ElementHandle} from '../api/ElementHandle.js'; import {assert} from '../util/assert.js'; import {CDPSession} from './Connection.js'; -import type {Frame} from './Frame.js'; -import type {WaitForSelectorOptions} from './IsolatedWorld.js'; import {IterableUtil} from './IterableUtil.js'; import {QueryHandler, QuerySelector} from './QueryHandler.js'; import {AwaitableIterable} from './types.js'; @@ -87,20 +85,16 @@ const parseARIASelector = (selector: string): ARIASelector => { return queryOptions; }; -/** - * @internal - */ -export interface ARIAQuerySelectorContext { - __ariaQuerySelector(node: Node, selector: string): Promise; -} - /** * @internal */ export class ARIAQueryHandler extends QueryHandler { - static override querySelector: QuerySelector = async (node, selector) => { - const context = globalThis as unknown as ARIAQuerySelectorContext; - return context.__ariaQuerySelector(node, selector); + static override querySelector: QuerySelector = async ( + node, + selector, + {ariaQuerySelector} + ) => { + return ariaQuerySelector(node, selector); }; static override async *queryAll( @@ -124,17 +118,4 @@ export class ARIAQueryHandler extends QueryHandler { ): Promise | null> => { return (await IterableUtil.first(this.queryAll(element, selector))) ?? null; }; - - static override async waitFor( - elementOrFrame: ElementHandle | Frame, - selector: string, - options: WaitForSelectorOptions - ): Promise | null> { - return super.waitFor( - elementOrFrame, - selector, - options, - new Map([['__ariaQuerySelector', this.queryOne]]) - ); - } } diff --git a/packages/puppeteer-core/src/common/Binding.ts b/packages/puppeteer-core/src/common/Binding.ts index fc99896b705..b9fdb0f40e1 100644 --- a/packages/puppeteer-core/src/common/Binding.ts +++ b/packages/puppeteer-core/src/common/Binding.ts @@ -6,16 +6,19 @@ import {debugError} from './util.js'; /** * @internal */ -export class Binding { +export class Binding { #name: string; - #fn: (...args: unknown[]) => unknown; - constructor(name: string, fn: (...args: unknown[]) => unknown) { + #fn: (...args: Params) => unknown; + constructor(name: string, fn: (...args: Params) => unknown) { this.#name = name; this.#fn = fn; } + get name(): string { + return this.#name; + } + /** - * * @param context - Context to run the binding in; the context should have * the binding added to it beforehand. * @param id - ID of the call. This should come from the CDP @@ -25,7 +28,7 @@ export class Binding { async run( context: ExecutionContext, id: number, - args: unknown[], + args: Params, isTrivial: boolean ): Promise { const garbage = []; diff --git a/packages/puppeteer-core/src/common/CustomQueryHandler.ts b/packages/puppeteer-core/src/common/CustomQueryHandler.ts index e5aec9c224f..2d4dccfb96d 100644 --- a/packages/puppeteer-core/src/common/CustomQueryHandler.ts +++ b/packages/puppeteer-core/src/common/CustomQueryHandler.ts @@ -14,13 +14,11 @@ * limitations under the License. */ -import {QueryHandler} from './QueryHandler.js'; -import {getQueryHandlerByName} from './GetQueryHandler.js'; - -/** - * @internal - */ -export const customQueryHandlers = new Map(); +import type PuppeteerUtil from '../injected/injected.js'; +import {assert} from '../util/assert.js'; +import {interpolateFunction, stringifyFunction} from '../util/Function.js'; +import {QueryHandler, QuerySelector, QuerySelectorAll} from './QueryHandler.js'; +import {scriptInjector} from './ScriptInjector.js'; /** * @public @@ -36,6 +34,154 @@ export interface CustomQueryHandler { queryAll?: (node: Node, selector: string) => Iterable; } +/** + * The registry of {@link CustomQueryHandler | custom query handlers}. + * + * @example + * + * ```ts + * Puppeteer.customQueryHandlers.register('lit', { … }); + * const aHandle = await page.$('lit/…'); + * ``` + * + * @internal + */ +export class CustomQueryHandlerRegistry { + #handlers = new Map< + string, + [registerScript: string, Handler: typeof QueryHandler] + >(); + + /** + * @internal + */ + get(name: string): typeof QueryHandler | undefined { + const handler = this.#handlers.get(name); + return handler ? handler[1] : undefined; + } + + /** + * Registers a {@link CustomQueryHandler | custom query handler}. + * + * @remarks + * After registration, the handler can be used everywhere where a selector is + * expected by prepending the selection string with `/`. The name is + * only allowed to consist of lower- and upper case latin letters. + * + * @example + * + * ```ts + * Puppeteer.customQueryHandlers.register('lit', { … }); + * const aHandle = await page.$('lit/…'); + * ``` + * + * @param name - Name to register under. + * @param queryHandler - {@link CustomQueryHandler | Custom query handler} to + * register. + * + * @internal + */ + register(name: string, handler: CustomQueryHandler): void { + if (this.#handlers.has(name)) { + throw new Error(`Cannot register over existing handler: ${name}`); + } + assert( + !this.#handlers.has(name), + `Cannot register over existing handler: ${name}` + ); + assert( + /^[a-zA-Z]+$/.test(name), + `Custom query handler names may only contain [a-zA-Z]` + ); + assert( + handler.queryAll || handler.queryOne, + `At least one query method must be implemented.` + ); + + const Handler = class extends QueryHandler { + static override querySelectorAll: QuerySelectorAll = interpolateFunction( + (node, selector, PuppeteerUtil) => { + return PuppeteerUtil.customQuerySelectors + .get(PLACEHOLDER('name'))! + .querySelectorAll(node, selector); + }, + {name: JSON.stringify(name)} + ); + static override querySelector: QuerySelector = interpolateFunction( + (node, selector, PuppeteerUtil) => { + return PuppeteerUtil.customQuerySelectors + .get(PLACEHOLDER('name'))! + .querySelector(node, selector); + }, + {name: JSON.stringify(name)} + ); + }; + const registerScript = interpolateFunction( + (PuppeteerUtil: PuppeteerUtil) => { + PuppeteerUtil.customQuerySelectors.register(PLACEHOLDER('name'), { + queryAll: PLACEHOLDER('queryAll'), + queryOne: PLACEHOLDER('queryOne'), + }); + }, + { + name: JSON.stringify(name), + queryAll: handler.queryAll + ? stringifyFunction(handler.queryAll) + : String(undefined), + queryOne: handler.queryOne + ? stringifyFunction(handler.queryOne) + : String(undefined), + } + ).toString(); + + this.#handlers.set(name, [registerScript, Handler]); + scriptInjector.append(registerScript); + } + + /** + * Unregisters the {@link CustomQueryHandler | custom query handler} for the + * given name. + * + * @throws `Error` if there is no handler under the given name. + * + * @internal + */ + unregister(name: string): void { + const handler = this.#handlers.get(name); + if (!handler) { + throw new Error(`Cannot unregister unknown handler: ${name}`); + } + scriptInjector.pop(handler[0]); + this.#handlers.delete(name); + } + + /** + * Gets the names of all {@link CustomQueryHandler | custom query handlers}. + * + * @internal + */ + names(): string[] { + return [...this.#handlers.keys()]; + } + + /** + * Unregisters all custom query handlers. + * + * @internal + */ + clear(): void { + for (const [registerScript] of this.#handlers) { + scriptInjector.pop(registerScript); + } + this.#handlers.clear(); + } +} + +/** + * @internal + */ +export const customQueryHandlers = new CustomQueryHandlerRegistry(); + /** * @deprecated Import {@link Puppeteer} and use the static method * {@link Puppeteer.registerCustomQueryHandler} @@ -46,22 +192,7 @@ export function registerCustomQueryHandler( name: string, handler: CustomQueryHandler ): void { - if (getQueryHandlerByName(name)) { - throw new Error(`A 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, - class extends QueryHandler { - static override querySelector = handler.queryOne; - static override querySelectorAll = handler.queryAll; - } - ); + customQueryHandlers.register(name, handler); } /** @@ -71,7 +202,7 @@ export function registerCustomQueryHandler( * @public */ export function unregisterCustomQueryHandler(name: string): void { - customQueryHandlers.delete(name); + customQueryHandlers.unregister(name); } /** @@ -81,7 +212,7 @@ export function unregisterCustomQueryHandler(name: string): void { * @public */ export function customQueryHandlerNames(): string[] { - return [...customQueryHandlers.keys()]; + return customQueryHandlers.names(); } /** diff --git a/packages/puppeteer-core/src/common/ElementHandle.ts b/packages/puppeteer-core/src/common/ElementHandle.ts index fee5c3769fb..284b00ca741 100644 --- a/packages/puppeteer-core/src/common/ElementHandle.ts +++ b/packages/puppeteer-core/src/common/ElementHandle.ts @@ -176,9 +176,9 @@ export class CDPElementHandle< override async $( selector: Selector ): Promise> | null> { - const {updatedSelector, queryHandler} = + const {updatedSelector, QueryHandler} = getQueryHandlerAndSelector(selector); - return (await queryHandler.queryOne( + return (await QueryHandler.queryOne( this, updatedSelector )) as CDPElementHandle> | null; @@ -187,10 +187,10 @@ export class CDPElementHandle< override async $$( selector: Selector ): Promise>>> { - const {updatedSelector, queryHandler} = + const {updatedSelector, QueryHandler} = getQueryHandlerAndSelector(selector); return IterableUtil.collect( - queryHandler.queryAll(this, updatedSelector) + QueryHandler.queryAll(this, updatedSelector) ) as Promise>>>; } @@ -256,9 +256,9 @@ export class CDPElementHandle< selector: Selector, options: WaitForSelectorOptions = {} ): Promise> | null> { - const {updatedSelector, queryHandler} = + const {updatedSelector, QueryHandler} = getQueryHandlerAndSelector(selector); - return (await queryHandler.waitFor( + return (await QueryHandler.waitFor( this, updatedSelector, options diff --git a/packages/puppeteer-core/src/common/ExecutionContext.ts b/packages/puppeteer-core/src/common/ExecutionContext.ts index 0a41584b70c..0cfa8954827 100644 --- a/packages/puppeteer-core/src/common/ExecutionContext.ts +++ b/packages/puppeteer-core/src/common/ExecutionContext.ts @@ -15,18 +15,20 @@ */ import {Protocol} from 'devtools-protocol'; -import {source as injectedSource} from '../generated/injected.js'; +import {JSHandle} from '../api/JSHandle.js'; import type PuppeteerUtil from '../injected/injected.js'; +import {stringifyFunction} from '../util/Function.js'; +import {ARIAQueryHandler} from './AriaQueryHandler.js'; +import {Binding} from './Binding.js'; import {CDPSession} from './Connection.js'; import {IsolatedWorld} from './IsolatedWorld.js'; -import {JSHandle} from '../api/JSHandle.js'; import {LazyArg} from './LazyArg.js'; +import {scriptInjector} from './ScriptInjector.js'; import {EvaluateFunc, HandleFor} from './types.js'; import { createJSHandle, getExceptionMessage, isString, - stringifyFunction, valueFromRemoteObject, } from './util.js'; import {CDPJSHandle} from './JSHandle.js'; @@ -94,16 +96,35 @@ export class ExecutionContext { #puppeteerUtil?: Promise>; get puppeteerUtil(): Promise> { - if (!this.#puppeteerUtil) { - this.#puppeteerUtil = this.evaluateHandle( - `(() => { - const module = {}; - ${injectedSource} - return module.exports.default; - })()` - ) as Promise>; + scriptInjector.inject(script => { + if (this.#puppeteerUtil) { + this.#puppeteerUtil.then(handle => { + handle.dispose(); + }); + } + this.#puppeteerUtil = Promise.all([ + this.#installGlobalBinding( + new Binding('__ariaQuerySelector', ARIAQueryHandler.queryOne) + ), + this.evaluateHandle(script) as Promise>, + ]).then(([, util]) => { + return util; + }); + }, !this.#puppeteerUtil); + return this.#puppeteerUtil as Promise>; + } + + async #installGlobalBinding(binding: Binding) { + try { + if (this._world) { + this._world._bindings.set(binding.name, binding); + await this._world._addBindingToContext(this, binding.name); + } + } catch { + // If the binding cannot be added, then either the browser doesn't support + // bindings (e.g. Firefox) or the context is broken. Either breakage is + // okay, so we ignore the error. } - return this.#puppeteerUtil; } /** @@ -272,8 +293,7 @@ export class ExecutionContext { let callFunctionOnPromise; try { callFunctionOnPromise = this._client.send('Runtime.callFunctionOn', { - functionDeclaration: - stringifyFunction(pageFunction) + '\n' + suffix + '\n', + functionDeclaration: `${stringifyFunction(pageFunction)}\n${suffix}\n`, executionContextId: this._contextId, arguments: await Promise.all(args.map(convertArgument.bind(this))), returnByValue, diff --git a/packages/puppeteer-core/src/common/Frame.ts b/packages/puppeteer-core/src/common/Frame.ts index 41e4bef141f..f244789b6c3 100644 --- a/packages/puppeteer-core/src/common/Frame.ts +++ b/packages/puppeteer-core/src/common/Frame.ts @@ -617,9 +617,9 @@ export class Frame { selector: Selector, options: WaitForSelectorOptions = {} ): Promise> | null> { - const {updatedSelector, queryHandler} = + const {updatedSelector, QueryHandler} = getQueryHandlerAndSelector(selector); - return (await queryHandler.waitFor( + return (await QueryHandler.waitFor( this, updatedSelector, options diff --git a/packages/puppeteer-core/src/common/GetQueryHandler.ts b/packages/puppeteer-core/src/common/GetQueryHandler.ts index 87745ac94dc..931f0ddbf1c 100644 --- a/packages/puppeteer-core/src/common/GetQueryHandler.ts +++ b/packages/puppeteer-core/src/common/GetQueryHandler.ts @@ -48,21 +48,23 @@ export function getQueryHandlerByName( */ export function getQueryHandlerAndSelector(selector: string): { updatedSelector: string; - queryHandler: typeof QueryHandler; + QueryHandler: typeof QueryHandler; } { for (const handlerMap of [ - customQueryHandlers, + customQueryHandlers.names().map(name => { + return [name, customQueryHandlers.get(name)!] as const; + }), Object.entries(BUILTIN_QUERY_HANDLERS), ]) { - for (const [name, queryHandler] of handlerMap) { + for (const [name, QueryHandler] of handlerMap) { for (const separator of QUERY_SEPARATORS) { const prefix = `${name}${separator}`; if (selector.startsWith(prefix)) { selector = selector.slice(prefix.length); - return {updatedSelector: selector, queryHandler}; + return {updatedSelector: selector, QueryHandler}; } } } } - return {updatedSelector: selector, queryHandler: CSSQueryHandler}; + return {updatedSelector: selector, QueryHandler: CSSQueryHandler}; } diff --git a/packages/puppeteer-core/src/common/IsolatedWorld.ts b/packages/puppeteer-core/src/common/IsolatedWorld.ts index ee5b56c7b9a..88692b5274e 100644 --- a/packages/puppeteer-core/src/common/IsolatedWorld.ts +++ b/packages/puppeteer-core/src/common/IsolatedWorld.ts @@ -41,6 +41,7 @@ import {TaskManager, WaitTask} from './WaitTask.js'; import type {ElementHandle} from '../api/ElementHandle.js'; import {Binding} from './Binding.js'; import {LazyArg} from './LazyArg.js'; +import {stringifyFunction} from '../util/Function.js'; /** * @public @@ -455,7 +456,7 @@ export class IsolatedWorld { LazyArg.create(context => { return context.puppeteerUtil; }), - queryOne.toString(), + stringifyFunction(queryOne as (...args: unknown[]) => unknown), selector, root, waitForVisible ? true : waitForHidden ? false : undefined diff --git a/packages/puppeteer-core/src/common/Puppeteer.ts b/packages/puppeteer-core/src/common/Puppeteer.ts index f5e46c58a97..0938c8c4ad0 100644 --- a/packages/puppeteer-core/src/common/Puppeteer.ts +++ b/packages/puppeteer-core/src/common/Puppeteer.ts @@ -20,13 +20,7 @@ import { _connectToCDPBrowser, } from './BrowserConnector.js'; import {ConnectionTransport} from './ConnectionTransport.js'; -import { - clearCustomQueryHandlers, - CustomQueryHandler, - customQueryHandlerNames, - registerCustomQueryHandler, - unregisterCustomQueryHandler, -} from './CustomQueryHandler.js'; +import {CustomQueryHandler, customQueryHandlers} from './CustomQueryHandler.js'; /** * Settings that are common to the Puppeteer class, regardless of environment. @@ -58,9 +52,18 @@ export interface ConnectOptions extends BrowserConnectOptions { * instance of {@link PuppeteerNode} when you import or require `puppeteer`. * That class extends `Puppeteer`, so has all the methods documented below as * well as all that are defined on {@link PuppeteerNode}. + * * @public */ export class Puppeteer { + /** + * Operations for {@link CustomQueryHandler | custom query handlers}. See + * {@link CustomQueryHandlerRegistry}. + * + * @internal + */ + static customQueryHandlers = customQueryHandlers; + /** * Registers a {@link CustomQueryHandler | custom query handler}. * @@ -87,28 +90,28 @@ export class Puppeteer { name: string, queryHandler: CustomQueryHandler ): void { - return registerCustomQueryHandler(name, queryHandler); + return this.customQueryHandlers.register(name, queryHandler); } /** * Unregisters a custom query handler for a given name. */ static unregisterCustomQueryHandler(name: string): void { - return unregisterCustomQueryHandler(name); + return this.customQueryHandlers.unregister(name); } /** * Gets the names of all custom query handlers. */ static customQueryHandlerNames(): string[] { - return customQueryHandlerNames(); + return this.customQueryHandlers.names(); } /** * Unregisters all custom query handlers. */ static clearCustomQueryHandlers(): void { - return clearCustomQueryHandlers(); + return this.customQueryHandlers.clear(); } /** diff --git a/packages/puppeteer-core/src/common/QueryHandler.ts b/packages/puppeteer-core/src/common/QueryHandler.ts index 03d68203edb..4238a9d479b 100644 --- a/packages/puppeteer-core/src/common/QueryHandler.ts +++ b/packages/puppeteer-core/src/common/QueryHandler.ts @@ -17,9 +17,9 @@ import {ElementHandle} from '../api/ElementHandle.js'; import type PuppeteerUtil from '../injected/injected.js'; import {assert} from '../util/assert.js'; -import {createFunction} from '../util/Function.js'; -import {transposeIterableHandle} from './HandleIterator.js'; +import {interpolateFunction, stringifyFunction} from '../util/Function.js'; import type {Frame} from './Frame.js'; +import {transposeIterableHandle} from './HandleIterator.js'; import type {WaitForSelectorOptions} from './IsolatedWorld.js'; import {MAIN_WORLD, PUPPETEER_WORLD} from './IsolatedWorlds.js'; import {LazyArg} from './LazyArg.js'; @@ -56,28 +56,23 @@ export class QueryHandler { return this.querySelector; } if (!this.querySelectorAll) { - throw new Error('Cannot create default query selector'); + throw new Error('Cannot create default `querySelector`.'); } - const querySelector: QuerySelector = async ( - node, - selector, - PuppeteerUtil - ) => { - const querySelectorAll = - 'FUNCTION_DEFINITION' as unknown as QuerySelectorAll; - const results = querySelectorAll(node, selector, PuppeteerUtil); - for await (const result of results) { - return result; + return (this.querySelector = interpolateFunction( + async (node, selector, PuppeteerUtil) => { + const querySelectorAll: QuerySelectorAll = + PLACEHOLDER('querySelectorAll'); + const results = querySelectorAll(node, selector, PuppeteerUtil); + for await (const result of results) { + return result; + } + return null; + }, + { + querySelectorAll: stringifyFunction(this.querySelectorAll), } - return null; - }; - - return (this.querySelector = createFunction( - querySelector - .toString() - .replace("'FUNCTION_DEFINITION'", this.querySelectorAll.toString()) - ) as typeof querySelector); + )); } static get _querySelectorAll(): QuerySelectorAll { @@ -85,26 +80,21 @@ export class QueryHandler { return this.querySelectorAll; } if (!this.querySelector) { - throw new Error('Cannot create default query selector'); + throw new Error('Cannot create default `querySelectorAll`.'); } - const querySelectorAll: QuerySelectorAll = async function* ( - node, - selector, - PuppeteerUtil - ) { - const querySelector = 'FUNCTION_DEFINITION' as unknown as QuerySelector; - const result = await querySelector(node, selector, PuppeteerUtil); - if (result) { - yield result; + return (this.querySelectorAll = interpolateFunction( + async function* (node, selector, PuppeteerUtil) { + const querySelector: QuerySelector = PLACEHOLDER('querySelector'); + const result = await querySelector(node, selector, PuppeteerUtil); + if (result) { + yield result; + } + }, + { + querySelector: stringifyFunction(this.querySelector), } - }; - - return (this.querySelectorAll = createFunction( - querySelectorAll - .toString() - .replace("'FUNCTION_DEFINITION'", this.querySelector.toString()) - ) as typeof querySelectorAll); + )); } /** @@ -160,8 +150,7 @@ export class QueryHandler { static async waitFor( elementOrFrame: ElementHandle | Frame, selector: string, - options: WaitForSelectorOptions, - bindings = new Map unknown>() + options: WaitForSelectorOptions ): Promise | null> { let frame: Frame; let element: ElementHandle | undefined; @@ -175,8 +164,7 @@ export class QueryHandler { this._querySelector, element, selector, - options, - bindings + options ); if (element) { await element.dispose(); diff --git a/packages/puppeteer-core/src/common/ScriptInjector.ts b/packages/puppeteer-core/src/common/ScriptInjector.ts new file mode 100644 index 00000000000..cb0c0395308 --- /dev/null +++ b/packages/puppeteer-core/src/common/ScriptInjector.ts @@ -0,0 +1,49 @@ +import {source as injectedSource} from '../generated/injected.js'; + +class ScriptInjector { + #updated = false; + #amendments = new Set(); + + // Appends a statement of the form `(PuppeteerUtil) => {...}`. + append(statement: string): void { + this.#update(() => { + this.#amendments.add(statement); + }); + } + + pop(statement: string): void { + this.#update(() => { + this.#amendments.delete(statement); + }); + } + + inject(inject: (script: string) => void, force = false) { + if (this.#updated || force) { + inject(this.#get()); + } + this.#updated = false; + } + + #update(callback: () => void): void { + callback(); + this.#updated = true; + } + + #get(): string { + return `(() => { + const module = {}; + ${injectedSource} + ${[...this.#amendments] + .map(statement => { + return `(${statement})(module.exports.default);`; + }) + .join('')} + return module.exports.default; + })()`; + } +} + +/** + * @internal + */ +export const scriptInjector = new ScriptInjector(); diff --git a/packages/puppeteer-core/src/common/WaitTask.ts b/packages/puppeteer-core/src/common/WaitTask.ts index 29aee81ed49..7058a9b869f 100644 --- a/packages/puppeteer-core/src/common/WaitTask.ts +++ b/packages/puppeteer-core/src/common/WaitTask.ts @@ -18,6 +18,7 @@ import {ElementHandle} from '../api/ElementHandle.js'; import {JSHandle} from '../api/JSHandle.js'; import type {Poller} from '../injected/Poller.js'; import {createDeferredPromise} from '../util/DeferredPromise.js'; +import {stringifyFunction} from '../util/Function.js'; import {Binding} from './Binding.js'; import {TimeoutError} from './Errors.js'; import {IsolatedWorld} from './IsolatedWorld.js'; @@ -68,7 +69,7 @@ export class WaitTask { this.#fn = `() => {return (${fn});}`; break; default: - this.#fn = fn.toString(); + this.#fn = stringifyFunction(fn); break; } this.#args = args; diff --git a/packages/puppeteer-core/src/common/bidi/Page.ts b/packages/puppeteer-core/src/common/bidi/Page.ts index dd82a701e48..00bb1fc4603 100644 --- a/packages/puppeteer-core/src/common/bidi/Page.ts +++ b/packages/puppeteer-core/src/common/bidi/Page.ts @@ -17,10 +17,11 @@ import {Page as PageBase} from '../../api/Page.js'; import {Connection} from './Connection.js'; import type {EvaluateFunc, HandleFor} from '../types.js'; -import {isString, stringifyFunction} from '../util.js'; +import {isString} from '../util.js'; import {BidiSerializer} from './Serializer.js'; import {JSHandle} from './JSHandle.js'; import {Reference} from './types.js'; +import {stringifyFunction} from '../../util/Function.js'; /** * @internal diff --git a/packages/puppeteer-core/src/common/util.ts b/packages/puppeteer-core/src/common/util.ts index 1e2180698f2..3eb0476d5cf 100644 --- a/packages/puppeteer-core/src/common/util.ts +++ b/packages/puppeteer-core/src/common/util.ts @@ -438,29 +438,3 @@ export async function getReadableFromProtocolStream( }, }); } - -/** - * @internal - */ -export function stringifyFunction(expression: Function): string { - let functionText = expression.toString(); - try { - new Function('(' + functionText + ')'); - } catch (error) { - // This means we might have a function shorthand. Try another - // time prefixing 'function '. - if (functionText.startsWith('async ')) { - functionText = - 'async function ' + functionText.substring('async '.length); - } else { - functionText = 'function ' + functionText; - } - try { - new Function('(' + functionText + ')'); - } catch (error) { - // We tried hard to serialize, but there's a weird beast here. - throw new Error('Passed function is not well-serializable!'); - } - } - return functionText; -} diff --git a/packages/puppeteer-core/src/injected/ARIAQuerySelector.ts b/packages/puppeteer-core/src/injected/ARIAQuerySelector.ts new file mode 100644 index 00000000000..52b3ea90565 --- /dev/null +++ b/packages/puppeteer-core/src/injected/ARIAQuerySelector.ts @@ -0,0 +1,31 @@ +/** + * Copyright 2022 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. + */ + +declare global { + interface Window { + /** + * @internal + */ + __ariaQuerySelector(root: Node, selector: string): Promise; + } +} + +export const ariaQuerySelector = ( + root: Node, + selector: string +): Promise => { + return window.__ariaQuerySelector(root, selector); +}; diff --git a/packages/puppeteer-core/src/injected/CustomQuerySelector.ts b/packages/puppeteer-core/src/injected/CustomQuerySelector.ts new file mode 100644 index 00000000000..869e385e63e --- /dev/null +++ b/packages/puppeteer-core/src/injected/CustomQuerySelector.ts @@ -0,0 +1,68 @@ +/** + * Copyright 2023 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. + */ + +import type {CustomQueryHandler} from '../common/CustomQueryHandler.js'; + +export interface CustomQuerySelector { + querySelector(root: Node, selector: string): Node | null; + querySelectorAll(root: Node, selector: string): Iterable; +} + +/** + * This class mimics the injected {@link CustomQuerySelectorRegistry}. + */ +class CustomQuerySelectorRegistry { + #selectors = new Map(); + + register(name: string, handler: CustomQueryHandler): void { + if (!handler.queryOne && handler.queryAll) { + const querySelectorAll = handler.queryAll; + handler.queryOne = (node, selector) => { + for (const result of querySelectorAll(node, selector)) { + return result; + } + return null; + }; + } else if (handler.queryOne && !handler.queryAll) { + const querySelector = handler.queryOne; + handler.queryAll = (node, selector) => { + const result = querySelector(node, selector); + return result ? [result] : []; + }; + } else if (!handler.queryOne || !handler.queryAll) { + throw new Error('At least one query method must be defined.'); + } + + this.#selectors.set(name, { + querySelector: handler.queryOne, + querySelectorAll: handler.queryAll!, + }); + } + + unregister(name: string): void { + this.#selectors.delete(name); + } + + get(name: string): CustomQuerySelector | undefined { + return this.#selectors.get(name); + } + + clear() { + this.#selectors.clear(); + } +} + +export const customQuerySelectors = new CustomQuerySelectorRegistry(); diff --git a/packages/puppeteer-core/src/injected/injected.ts b/packages/puppeteer-core/src/injected/injected.ts index 21d8cdeeb3c..f4000b92400 100644 --- a/packages/puppeteer-core/src/injected/injected.ts +++ b/packages/puppeteer-core/src/injected/injected.ts @@ -16,26 +16,30 @@ import {createDeferredPromise} from '../util/DeferredPromise.js'; import {createFunction} from '../util/Function.js'; -import {RAFPoller, MutationPoller, IntervalPoller} from './Poller.js'; +import * as ARIAQuerySelector from './ARIAQuerySelector.js'; +import * as CustomQuerySelectors from './CustomQuerySelector.js'; +import * as PierceQuerySelector from './PierceQuerySelector.js'; +import {IntervalPoller, MutationPoller, RAFPoller} from './Poller.js'; import { - isSuitableNodeForTextMatching, createTextContent, + isSuitableNodeForTextMatching, } from './TextContent.js'; import * as TextQuerySelector from './TextQuerySelector.js'; -import * as XPathQuerySelector from './XPathQuerySelector.js'; -import * as PierceQuerySelector from './PierceQuerySelector.js'; import * as util from './util.js'; +import * as XPathQuerySelector from './XPathQuerySelector.js'; /** * @internal */ const PuppeteerUtil = Object.freeze({ - ...util, - ...TextQuerySelector, - ...XPathQuerySelector, + ...ARIAQuerySelector, + ...CustomQuerySelectors, ...PierceQuerySelector, - createFunction, + ...TextQuerySelector, + ...util, + ...XPathQuerySelector, createDeferredPromise, + createFunction, createTextContent, IntervalPoller, isSuitableNodeForTextMatching, diff --git a/packages/puppeteer-core/src/util/Function.ts b/packages/puppeteer-core/src/util/Function.ts index dc33bf4d0bf..cdf09ba195c 100644 --- a/packages/puppeteer-core/src/util/Function.ts +++ b/packages/puppeteer-core/src/util/Function.ts @@ -33,3 +33,66 @@ export const createFunction = ( createdFunctions.set(functionValue, fn); return fn; }; + +/** + * @internal + */ +export function stringifyFunction(fn: (...args: never) => unknown): string { + let value = fn.toString(); + try { + new Function(`(${value})`); + } catch { + // This means we might have a function shorthand (e.g. `test(){}`). Let's + // try prefixing. + let prefix = 'function '; + if (value.startsWith('async ')) { + prefix = `async ${prefix}`; + value = value.substring('async '.length); + } + value = `${prefix}${value}`; + try { + new Function(`(${value})`); + } catch { + // We tried hard to serialize, but there's a weird beast here. + throw new Error('Passed function cannot be serialized!'); + } + } + return value; +} + +/** + * Replaces `PLACEHOLDER`s with the given replacements. + * + * All replacements must be valid JS code. + * + * @example + * + * ```ts + * interpolateFunction(() => PLACEHOLDER('test'), {test: 'void 0'}); + * // Equivalent to () => void 0 + * ``` + * + * @internal + */ +export const interpolateFunction = unknown>( + fn: T, + replacements: Record +): T => { + let value = stringifyFunction(fn); + for (const [name, jsValue] of Object.entries(replacements)) { + value = value.replace( + new RegExp(`PLACEHOLDER\\(\\s*(?:'${name}'|"${name}")\\s*\\)`, 'g'), + jsValue + ); + } + return createFunction(value) as unknown as T; +}; + +declare global { + /** + * Used for interpolation with {@link interpolateFunction}. + * + * @internal + */ + function PLACEHOLDER(name: string): T; +} diff --git a/test/src/elementhandle.spec.ts b/test/src/elementhandle.spec.ts index 20c7b8f564d..035a777c200 100644 --- a/test/src/elementhandle.spec.ts +++ b/test/src/elementhandle.spec.ts @@ -702,6 +702,27 @@ describe('ElementHandle specs', function () { }); expect(txtContents).toBe('textcontent'); }); + + it('should work with function shorthands', async () => { + const {page} = getTestState(); + await page.setContent('
'); + + Puppeteer.registerCustomQueryHandler('getById', { + // This is a function shorthand + queryOne(_element, selector) { + return document.querySelector(`[id="${selector}"]`); + }, + }); + + const element = (await page.$( + 'getById/foo' + )) as ElementHandle; + expect( + await page.evaluate(element => { + return element.id; + }, element) + ).toBe('foo'); + }); }); describe('Element.toElement', () => {