feat!: type inference for evaluation types

src/api-docs-entry.ts

@@ -72,7 +72,7 @@ export * from './common/Tracing.js';
 export * from './common/NetworkManager.js';
 export * from './common/WebWorker.js';
 export * from './common/USKeyboardLayout.js';
-export * from './common/EvalTypes.js';
+export * from './common/types.js';
 export * from './common/PDFOptions.js';
 export * from './common/TimeoutSettings.js';
 export * from './common/LifecycleWatcher.js';

src/common/AriaQueryHandler.ts

@@ -141,7 +141,7 @@ const queryAll = async (
 const queryAllArray = async (
   element: ElementHandle,
   selector: string
-): Promise<JSHandle> => {
+): Promise<JSHandle<Element[]>> => {
   const elementHandles = await queryAll(element, selector);
   const exeCtx = element.executionContext();
   const jsHandle = exeCtx.evaluateHandle((...elements) => {
@@ -153,7 +153,7 @@ const queryAllArray = async (
 /**
  * @internal
  */
-export const _ariaHandler: InternalQueryHandler = {
+export const ariaHandler: InternalQueryHandler = {
   queryOne,
   waitFor,
   queryAll,

src/common/Connection.ts

@@ -33,8 +33,8 @@ export {ConnectionTransport, ProtocolMapping};
  * @public
  */
 export interface ConnectionCallback {
-  resolve: Function;
-  reject: Function;
+  resolve(args: unknown): void;
+  reject(args: unknown): void;
   error: ProtocolError;
   method: string;
 }

src/common/DOMWorld.ts

@@ -18,28 +18,21 @@ import {Protocol} from 'devtools-protocol';
 import {assert} from './assert.js';
 import {CDPSession} from './Connection.js';
 import {TimeoutError} from './Errors.js';
-import {
-  EvaluateFn,
-  EvaluateFnReturnType,
-  EvaluateHandleFn,
-  SerializableOrJSHandle,
-  UnwrapPromiseLike,
-  WrapElementHandle,
-} from './EvalTypes.js';
 import {ExecutionContext} from './ExecutionContext.js';
 import {Frame, FrameManager} from './FrameManager.js';
+import {MouseButton} from './Input.js';
+import {ElementHandle, JSHandle} from './JSHandle.js';
+import {LifecycleWatcher, PuppeteerLifeCycleEvent} from './LifecycleWatcher.js';
+import {_getQueryHandlerAndSelector} from './QueryHandler.js';
+import {TimeoutSettings} from './TimeoutSettings.js';
+import {EvaluateFunc, EvaluateParams, HandleFor} from './types.js';
 import {
   debugError,
   isNumber,
   isString,
   makePredicateString,
   pageBindingInitString,
 } from './util.js';
-import {MouseButton} from './Input.js';
-import {ElementHandle, JSHandle} from './JSHandle.js';
-import {LifecycleWatcher, PuppeteerLifeCycleEvent} from './LifecycleWatcher.js';
-import {_getQueryHandlerAndSelector} from './QueryHandler.js';
-import {TimeoutSettings} from './TimeoutSettings.js';
 
 // predicateQueryHandler and checkWaitForOptions are declared here so that
 // TypeScript knows about them when used in the predicate function below.
@@ -184,30 +177,45 @@ export class DOMWorld {
     return this.#contextPromise;
   }
 
-  async evaluateHandle<HandlerType extends JSHandle = JSHandle>(
-    pageFunction: EvaluateHandleFn,
-    ...args: SerializableOrJSHandle[]
-  ): Promise<HandlerType> {
+  async evaluateHandle<
+    Params extends unknown[],
+    Func extends EvaluateFunc<Params> = EvaluateFunc<Params>
+  >(
+    pageFunction: Func | string,
+    ...args: EvaluateParams<Params>
+  ): Promise<HandleFor<Awaited<ReturnType<Func>>>> {
     const context = await this.executionContext();
     return context.evaluateHandle(pageFunction, ...args);
   }
 
-  async evaluate<T extends EvaluateFn>(
-    pageFunction: T,
-    ...args: SerializableOrJSHandle[]
-  ): Promise<UnwrapPromiseLike<EvaluateFnReturnType<T>>> {
+  async evaluate<
+    Params extends unknown[],
+    Func extends EvaluateFunc<Params> = EvaluateFunc<Params>
+  >(
+    pageFunction: Func | string,
+    ...args: EvaluateParams<Params>
+  ): Promise<Awaited<ReturnType<Func>>> {
     const context = await this.executionContext();
-    return context.evaluate<UnwrapPromiseLike<EvaluateFnReturnType<T>>>(
-      pageFunction,
-      ...args
-    );
+    return context.evaluate(pageFunction, ...args);
   }
 
-  async $<T extends Element = Element>(
-    selector: string
-  ): Promise<ElementHandle<T> | null> {
+  async $<Selector extends keyof HTMLElementTagNameMap>(
+    selector: Selector
+  ): Promise<ElementHandle<HTMLElementTagNameMap[Selector]> | null>;
+  async $(selector: string): Promise<ElementHandle | null>;
+  async $(selector: string): Promise<ElementHandle | null> {
     const document = await this._document();
-    const value = await document.$<T>(selector);
+    const value = await document.$(selector);
+    return value;
+  }
+
+  async $$<Selector extends keyof HTMLElementTagNameMap>(
+    selector: Selector
+  ): Promise<ElementHandle<HTMLElementTagNameMap[Selector]>[]>;
+  async $$(selector: string): Promise<ElementHandle[]>;
+  async $$(selector: string): Promise<ElementHandle[]> {
+    const document = await this._document();
+    const value = await document.$$(selector);
     return value;
   }
 
@@ -235,40 +243,74 @@ export class DOMWorld {
     return value;
   }
 
-  async $eval<ReturnType>(
+  async $eval<
+    Selector extends keyof HTMLElementTagNameMap,
+    Params extends unknown[],
+    Func extends EvaluateFunc<
+      [HTMLElementTagNameMap[Selector], ...Params]
+    > = EvaluateFunc<[HTMLElementTagNameMap[Selector], ...Params]>
+  >(
+    selector: Selector,
+    pageFunction: Func | string,
+    ...args: EvaluateParams<Params>
+  ): Promise<Awaited<ReturnType<Func>>>;
+  async $eval<
+    Params extends unknown[],
+    Func extends EvaluateFunc<[Element, ...Params]> = EvaluateFunc<
+      [Element, ...Params]
+    >
+  >(
     selector: string,
-    pageFunction: (
-      element: Element,
-      ...args: unknown[]
-    ) => ReturnType | Promise<ReturnType>,
-    ...args: SerializableOrJSHandle[]
-  ): Promise<WrapElementHandle<ReturnType>> {
-    const document = await this._document();
-    return document.$eval<ReturnType>(selector, pageFunction, ...args);
-  }
-
-  async $$eval<ReturnType>(
+    pageFunction: Func | string,
+    ...args: EvaluateParams<Params>
+  ): Promise<Awaited<ReturnType<Func>>>;
+  async $eval<
+    Params extends unknown[],
+    Func extends EvaluateFunc<[Element, ...Params]> = EvaluateFunc<
+      [Element, ...Params]
+    >
+  >(
     selector: string,
-    pageFunction: (
-      elements: Element[],
-      ...args: unknown[]
-    ) => ReturnType | Promise<ReturnType>,
-    ...args: SerializableOrJSHandle[]
-  ): Promise<WrapElementHandle<ReturnType>> {
+    pageFunction: Func | string,
+    ...args: EvaluateParams<Params>
+  ): Promise<Awaited<ReturnType<Func>>> {
     const document = await this._document();
-    const value = await document.$$eval<ReturnType>(
-      selector,
-      pageFunction,
-      ...args
-    );
-    return value;
+    return document.$eval(selector, pageFunction, ...args);
   }
 
-  async $$<T extends Element = Element>(
-    selector: string
-  ): Promise<Array<ElementHandle<T>>> {
+  async $$eval<
+    Selector extends keyof HTMLElementTagNameMap,
+    Params extends unknown[],
+    Func extends EvaluateFunc<
+      [HTMLElementTagNameMap[Selector][], ...Params]
+    > = EvaluateFunc<[HTMLElementTagNameMap[Selector][], ...Params]>
+  >(
+    selector: Selector,
+    pageFunction: Func | string,
+    ...args: EvaluateParams<Params>
+  ): Promise<Awaited<ReturnType<Func>>>;
+  async $$eval<
+    Params extends unknown[],
+    Func extends EvaluateFunc<[Element[], ...Params]> = EvaluateFunc<
+      [Element[], ...Params]
+    >
+  >(
+    selector: string,
+    pageFunction: Func | string,
+    ...args: EvaluateParams<Params>
+  ): Promise<Awaited<ReturnType<Func>>>;
+  async $$eval<
+    Params extends unknown[],
+    Func extends EvaluateFunc<[Element[], ...Params]> = EvaluateFunc<
+      [Element[], ...Params]
+    >
+  >(
+    selector: string,
+    pageFunction: Func | string,
+    ...args: EvaluateParams<Params>
+  ): Promise<Awaited<ReturnType<Func>>> {
     const document = await this._document();
-    const value = await document.$$<T>(selector);
+    const value = await document.$$eval(selector, pageFunction, ...args);
     return value;
   }
 
@@ -298,7 +340,7 @@ export class DOMWorld {
     } = options;
     // We rely upon the fact that document.open() will reset frame lifecycle with "init"
     // lifecycle event. @see https://crrev.com/608658
-    await this.evaluate<(x: string) => void>(html => {
+    await this.evaluate(html => {
       document.open();
       document.write(html);
       document.close();
@@ -536,7 +578,6 @@ export class DOMWorld {
 
     async function addStyleContent(content: string): Promise<HTMLElement> {
       const style = document.createElement('style');
-      style.type = 'text/css';
       style.appendChild(document.createTextNode(content));
       const promise = new Promise((res, rej) => {
         style.onload = res;
@@ -598,6 +639,14 @@ export class DOMWorld {
     await handle.dispose();
   }
 
+  async waitForSelector<Selector extends keyof HTMLElementTagNameMap>(
+    selector: Selector,
+    options: WaitForSelectorOptions
+  ): Promise<ElementHandle<HTMLElementTagNameMap[Selector]> | null>;
+  async waitForSelector(
+    selector: string,
+    options: WaitForSelectorOptions
+  ): Promise<ElementHandle | null>;
   async waitForSelector(
     selector: string,
     options: WaitForSelectorOptions
@@ -825,7 +874,7 @@ export class DOMWorld {
   waitForFunction(
     pageFunction: Function | string,
     options: {polling?: string | number; timeout?: number} = {},
-    ...args: SerializableOrJSHandle[]
+    ...args: unknown[]
   ): Promise<JSHandle> {
     const {polling = 'raf', timeout = this.#timeoutSettings.timeout()} =
       options;
@@ -860,7 +909,7 @@ export interface WaitTaskOptions {
   polling: string | number;
   timeout: number;
   binding?: PageBinding;
-  args: SerializableOrJSHandle[];
+  args: unknown[];
   root?: ElementHandle;
 }
 
@@ -871,11 +920,11 @@ const noop = (): void => {};
  */
 export class WaitTask {
   #domWorld: DOMWorld;
-  #polling: string | number;
+  #polling: 'raf' | 'mutation' | number;
   #timeout: number;
   #predicateBody: string;
   #predicateAcceptsContextElement: boolean;
-  #args: SerializableOrJSHandle[];
+  #args: unknown[];
   #binding?: PageBinding;
   #runCount = 0;
   #resolve: (x: JSHandle) => void = noop;