From 8f9eaf7cba53584436703bc7e35f27933c1a8c94 Mon Sep 17 00:00:00 2001
From: Yury Semikhatsky <yurys@chromium.org>
Date: Tue, 21 Dec 2021 13:17:45 -0800
Subject: [PATCH] fix(types): allow specifying element type in locator.evaluate
 callback (#11048)

---
 packages/playwright-core/types/types.d.ts |  8 +++---
 utils/generate_types/overrides.d.ts       |  8 +++---
 utils/generate_types/test/test.ts         | 30 +++++++++++++++++++++++
 3 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/packages/playwright-core/types/types.d.ts b/packages/playwright-core/types/types.d.ts
index b109101019..5605952251 100644
--- a/packages/playwright-core/types/types.d.ts
+++ b/packages/playwright-core/types/types.d.ts
@@ -8508,7 +8508,7 @@ export interface Locator {
    * @param arg Optional argument to pass to `pageFunction`.
    * @param options
    */
-  evaluate<R, Arg>(pageFunction: PageFunctionOn<SVGElement | HTMLElement, Arg, R>, arg: Arg, options?: {
+  evaluate<R, Arg, E extends SVGElement | HTMLElement = SVGElement | HTMLElement>(pageFunction: PageFunctionOn<E, Arg, R>, arg: Arg, options?: {
     timeout?: number;
   }): Promise<R>;
   /**
@@ -8530,7 +8530,7 @@ export interface Locator {
    * @param arg Optional argument to pass to `pageFunction`.
    * @param options
    */
-  evaluate<R>(pageFunction: PageFunctionOn<SVGElement | HTMLElement, void, R>, options?: {
+  evaluate<R, E extends SVGElement | HTMLElement = SVGElement | HTMLElement>(pageFunction: PageFunctionOn<E, void, R>, options?: {
     timeout?: number;
   }): Promise<R>;
   /**
@@ -8551,7 +8551,7 @@ export interface Locator {
    * @param pageFunction Function to be evaluated in the page context.
    * @param arg Optional argument to pass to `pageFunction`.
    */
-  evaluateAll<R, Arg>(pageFunction: PageFunctionOn<(SVGElement | HTMLElement)[], Arg, R>, arg: Arg): Promise<R>;
+  evaluateAll<R, Arg, E extends SVGElement | HTMLElement = SVGElement | HTMLElement>(pageFunction: PageFunctionOn<E[], Arg, R>, arg: Arg): Promise<R>;
   /**
    * The method finds all elements matching the specified locator and passes an array of matched elements as a first argument
    * to `pageFunction`. Returns the result of `pageFunction` invocation.
@@ -8570,7 +8570,7 @@ export interface Locator {
    * @param pageFunction Function to be evaluated in the page context.
    * @param arg Optional argument to pass to `pageFunction`.
    */
-  evaluateAll<R>(pageFunction: PageFunctionOn<(SVGElement | HTMLElement)[], void, R>): Promise<R>;
+  evaluateAll<R, E extends SVGElement | HTMLElement = SVGElement | HTMLElement>(pageFunction: PageFunctionOn<E[], void, R>): Promise<R>;
   /**
    * Resolves given locator to the first matching DOM element. If no elements matching the query are visible, waits for them
    * up to a given timeout. If multiple elements match the selector, throws.
diff --git a/utils/generate_types/overrides.d.ts b/utils/generate_types/overrides.d.ts
index 6489f7826a..9a4cc72041 100644
--- a/utils/generate_types/overrides.d.ts
+++ b/utils/generate_types/overrides.d.ts
@@ -142,14 +142,14 @@ export interface ElementHandle<T=Node> extends JSHandle<T> {
 }
 
 export interface Locator {
-  evaluate<R, Arg>(pageFunction: PageFunctionOn<SVGElement | HTMLElement, Arg, R>, arg: Arg, options?: {
+  evaluate<R, Arg, E extends SVGElement | HTMLElement = SVGElement | HTMLElement>(pageFunction: PageFunctionOn<E, Arg, R>, arg: Arg, options?: {
     timeout?: number;
   }): Promise<R>;
-  evaluate<R>(pageFunction: PageFunctionOn<SVGElement | HTMLElement, void, R>, options?: {
+  evaluate<R, E extends SVGElement | HTMLElement = SVGElement | HTMLElement>(pageFunction: PageFunctionOn<E, void, R>, options?: {
     timeout?: number;
   }): Promise<R>;
-  evaluateAll<R, Arg>(pageFunction: PageFunctionOn<(SVGElement | HTMLElement)[], Arg, R>, arg: Arg): Promise<R>;
-  evaluateAll<R>(pageFunction: PageFunctionOn<(SVGElement | HTMLElement)[], void, R>): Promise<R>;
+  evaluateAll<R, Arg, E extends SVGElement | HTMLElement = SVGElement | HTMLElement>(pageFunction: PageFunctionOn<E[], Arg, R>, arg: Arg): Promise<R>;
+  evaluateAll<R, E extends SVGElement | HTMLElement = SVGElement | HTMLElement>(pageFunction: PageFunctionOn<E[], void, R>): Promise<R>;
   elementHandle(options?: {
     timeout?: number;
   }): Promise<null|ElementHandle<SVGElement | HTMLElement>>;
diff --git a/utils/generate_types/test/test.ts b/utils/generate_types/test/test.ts
index b34f1fa91e..0f26e3c8fe 100644
--- a/utils/generate_types/test/test.ts
+++ b/utils/generate_types/test/test.ts
@@ -387,6 +387,36 @@ playwright.chromium.launch().then(async browser => {
   await browser.close();
 })();
 
+// test locator.evaluate
+(async () => {
+  const browser = await playwright.firefox.launch();
+  const page = await browser.newPage();
+  const locator = page.locator('.foo');
+  {
+    const result = await locator.evaluate((sel: HTMLSelectElement) => sel.options[sel.selectedIndex].textContent)
+    const assertion: AssertType<string, typeof result> = true;
+  }
+  {
+    const result = await locator.evaluate((media: HTMLMediaElement, dummy) => media.duration, 10);
+    const assertion: AssertType<number, typeof result> = true;
+  }
+  {
+    await locator.evaluate((input: HTMLInputElement) => {})
+  }
+  {
+    const list = await locator.evaluateAll((i: HTMLInputElement[]) => i.length);
+    const assertion: AssertType<number, typeof list> = true;
+  }
+  {
+    const list = await locator.evaluateAll((i: HTMLInputElement[], dummy) => i.length, 10);
+    const assertion: AssertType<number, typeof list> = true;
+  }
+  {
+    await locator.evaluateAll((sel: HTMLSelectElement[]) => {})
+  }
+  await browser.close();
+})();
+
 // waitForEvent
 (async () => {
   const browser = await playwright.webkit.launch();