Merge pull request #32759 from bgschiller/puppeteer-evaluate-return-types

puppeteer: preserve type information in Page.evaluate and friends

types/puppeteer/index.d.ts

@@ -507,6 +507,7 @@ export interface EmulateOptions {
 }
 
 export type EvaluateFn = string | ((...args: any[]) => any);
+export type EvaluateFnReturnType<T extends EvaluateFn> = T extends ((...args: any[]) => infer R) ? R : any;
 
 export type LoadEvent =
   | "load"
@@ -746,10 +747,10 @@ export interface Worker {
    * If the function passed to the `worker.evaluate` returns a non-Serializable value,
    * then `worker.evaluate` resolves to `undefined`.
    */
-  evaluate<T>(
-    pageFunction: (...args: any[]) => T | Promise<T>,
+  evaluate<T extends EvaluateFn>(
+    pageFunction: T,
     ...args: SerializableOrJSHandle[],
-  ): Promise<T>;
+  ): Promise<EvaluateFnReturnType<T>>;
 
   /**
    * The only difference between `worker.evaluate` and `worker.evaluateHandle` is
@@ -858,10 +859,10 @@ export interface ElementHandle<E extends Element = Element> extends JSHandle, Ev
 
 /** The class represents a context for JavaScript execution. */
 export interface ExecutionContext {
-  evaluate(
-    fn: EvaluateFn,
+  evaluate<F extends EvaluateFn>(
+    fn: F,
     ...args: SerializableOrJSHandle[]
-  ): Promise<any>;
+  ): Promise<EvaluateFnReturnType<F>>;
   evaluateHandle(
     fn: EvaluateFn,
     ...args: SerializableOrJSHandle[]
@@ -1173,10 +1174,10 @@ export interface FrameBase extends Evalable {
    * @param fn Function to be evaluated in browser context
    * @param args Arguments to pass to `fn`
    */
-  evaluate(
-    fn: EvaluateFn,
+  evaluate<F extends EvaluateFn>(
+    fn: F,
     ...args: SerializableOrJSHandle[]
-  ): Promise<any>;
+  ): Promise<EvaluateFnReturnType<F>>;
 
   /**
    * Evaluates a function in the page context.
@@ -1255,19 +1256,19 @@ export interface FrameBase extends Evalable {
    * Shortcut for waitForFunction.
    */
   waitFor(
-    selector: ((...args: any[]) => any) | string,
+    selector: EvaluateFn,
     options?: WaitForSelectorOptions,
     ...args: SerializableOrJSHandle[]
-  ): Promise<any>;
+  ): Promise<JSHandle>;
 
   /**
    * Allows waiting for various conditions.
    */
   waitForFunction(
-    fn: string | ((...args: any[]) => any),
+    fn: EvaluateFn,
     options?: PageFnOptions,
     ...args: SerializableOrJSHandle[]
-  ): Promise<any>;
+  ): Promise<JSHandle>;
 
   /**
    * Wait for the page navigation occur.

types/puppeteer/puppeteer-tests.ts

@@ -531,6 +531,35 @@ puppeteer.launch().then(async browser => {
   });
 });
 
+// evaluates return type of inner function
+(async () => {
+  const browser = await puppeteer.launch();
+  const page = await browser.newPage();
+  const s = await page.evaluate(() => document.body.innerHTML);
+  console.log('body html has length', s.length);
+});
+
+// even through a double promise.
+(async () => {
+  const browser = await puppeteer.launch();
+  const page = await browser.newPage();
+  const s = await page.evaluate(() => Promise.resolve(document.body.innerHTML));
+  console.log('body html has length', s.length);
+});
+
+// JSHandle.jsonValue produces compatible type
+(async () => {
+  const browser = await puppeteer.launch();
+  const page = await browser.newPage();
+  const s = await page
+    .waitForFunction(
+      (searchStrs: string[]) => searchStrs.find(v => document.body.innerText.includes(v)),
+      { timeout: 2000 },
+      ['once', 'upon', 'a', 'midnight', 'dreary'])
+    .then(j => j.jsonValue());
+  console.log('found in page', s.toLowerCase());
+});
+
 // Element access
 (async () => {
   const browser = await puppeteer.launch();
@@ -547,3 +576,45 @@ puppeteer.launch().then(async browser => {
     a: '1'
   });
 });
+
+// ElementHandles are well-typed
+(async () => {
+  const browser = await puppeteer.launch();
+  const page = await browser.newPage();
+  const link: puppeteer.JSHandle = await page.evaluateHandle(
+    () => document.body.querySelector('a')
+  );
+  const linkEl: puppeteer.ElementHandle | null = link.asElement();
+  if (linkEl !== null) {
+    const href = await page.evaluate(
+      (el: HTMLElement): string | null => el.getAttribute('href'),
+      linkEl);
+    console.log('href is', href);
+  }
+});
+
+// test $$eval return type
+(async () => {
+  const browser = await puppeteer.launch();
+  const page = await browser.newPage();
+
+  const paragraphContents: string[] = await page.$$eval(
+    'p', (ps: Element[]): string[] => ps.map(p => p.textContent || ''));
+  console.log('pgraph contents', paragraphContents);
+});
+
+// JSHandle of non-serializable works
+(async () => {
+  const browser = await puppeteer.launch();
+  const page = await browser.newPage();
+
+  const reHandle: puppeteer.JSHandle = await page.evaluateHandle(
+    () => /\s*bananas?\s*/i,
+  );
+  const numMatchingEls: number = await page.$$eval(
+    'p', (els: Element[], re: RegExp) =>
+      els.filter(el => el.textContent && re.test(el.textContent)).length,
+      reHandle
+  );
+  console.log('there are', numMatchingEls, 'banana paragaphs');
+});