From a9289fe0c03e0b21ed642440617925e49090dbfe Mon Sep 17 00:00:00 2001
From: Randolf <jrandolf@google.com>
Date: Wed, 9 Aug 2023 16:06:09 +0200
Subject: [PATCH] chore: implement `mapHandle` and `filterHandle`

---
 .../src/api/locators/FilteredLocator.ts       | 22 +++++++-----
 .../src/api/locators/Locator.ts               | 35 ++++++++++++++++++-
 .../src/api/locators/MappedLocator.ts         | 15 +++++---
 3 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/packages/puppeteer-core/src/api/locators/FilteredLocator.ts b/packages/puppeteer-core/src/api/locators/FilteredLocator.ts
index ca2a9da5b88e6..9f999fa8c3c3b 100644
--- a/packages/puppeteer-core/src/api/locators/FilteredLocator.ts
+++ b/packages/puppeteer-core/src/api/locators/FilteredLocator.ts
@@ -16,13 +16,13 @@
 
 import {
   Observable,
+  filter,
   from,
   map,
   mergeMap,
   throwIfEmpty,
 } from '../../../third_party/rxjs/rxjs.js';
 import {Awaitable, HandleFor} from '../../common/common.js';
-import {ElementHandle} from '../ElementHandle.js';
 
 import {DelegatedLocator} from './DelegatedLocator.js';
 import {ActionOptions, Locator} from './locators.js';
@@ -34,6 +34,13 @@ export type Predicate<From, To extends From = From> =
   | ((value: From) => value is To)
   | ((value: From) => Awaitable<boolean>);
 
+/**
+ * @internal
+ */
+export type HandlePredicate<From, To extends From = From> =
+  | ((value: HandleFor<From>, signal?: AbortSignal) => value is HandleFor<To>)
+  | ((value: HandleFor<From>, signal?: AbortSignal) => Awaitable<boolean>);
+
 /**
  * @internal
  */
@@ -41,9 +48,9 @@ export class FilteredLocator<From, To extends From> extends DelegatedLocator<
   From,
   To
 > {
-  #predicate: Predicate<From, To>;
+  #predicate: HandlePredicate<From, To>;
 
-  constructor(base: Locator<From>, predicate: Predicate<From, To>) {
+  constructor(base: Locator<From>, predicate: HandlePredicate<From, To>) {
     super(base);
     this.#predicate = predicate;
   }
@@ -59,12 +66,11 @@ export class FilteredLocator<From, To extends From> extends DelegatedLocator<
     return this.delegate._wait(options).pipe(
       mergeMap(handle => {
         return from(
-          (handle as ElementHandle<Node>).frame.waitForFunction(
-            this.#predicate,
-            {signal: options?.signal, timeout: this._timeout},
-            handle
-          )
+          Promise.resolve(this.#predicate(handle, options?.signal))
         ).pipe(
+          filter(value => {
+            return value;
+          }),
           map(() => {
             // SAFETY: It passed the predicate, so this is correct.
             return handle as HandleFor<To>;
diff --git a/packages/puppeteer-core/src/api/locators/Locator.ts b/packages/puppeteer-core/src/api/locators/Locator.ts
index 19a769c7f0c55..ede9449877e6f 100644
--- a/packages/puppeteer-core/src/api/locators/Locator.ts
+++ b/packages/puppeteer-core/src/api/locators/Locator.ts
@@ -48,6 +48,7 @@ import {
   Action,
   AwaitedLocator,
   FilteredLocator,
+  HandleMapper,
   MappedLocator,
   Mapper,
   Predicate,
@@ -702,7 +703,10 @@ export abstract class Locator<T> extends EventEmitter {
    * @public
    */
   map<To>(mapper: Mapper<T, To>): Locator<To> {
-    return new MappedLocator(this._clone(), mapper);
+    return new MappedLocator(this._clone(), handle => {
+      // SAFETY: TypeScript cannot deduce the type.
+      return (handle as any).evaluateHandle(mapper);
+    });
   }
 
   /**
@@ -713,9 +717,38 @@ export abstract class Locator<T> extends EventEmitter {
    * @public
    */
   filter<S extends T>(predicate: Predicate<T, S>): Locator<S> {
+    return new FilteredLocator(this._clone(), async (handle, signal) => {
+      await (handle as ElementHandle<Node>).frame.waitForFunction(
+        predicate,
+        {signal, timeout: this._timeout},
+        handle
+      );
+      return true;
+    });
+  }
+
+  /**
+   * Creates an expectation that is evaluated against located handles.
+   *
+   * If the expectations do not match, then the locator will retry.
+   *
+   * @internal
+   */
+  filterHandle<S extends T>(
+    predicate: Predicate<HandleFor<T>, HandleFor<S>>
+  ): Locator<S> {
     return new FilteredLocator(this._clone(), predicate);
   }
 
+  /**
+   * Maps the locator using the provided mapper.
+   *
+   * @internal
+   */
+  mapHandle<To>(mapper: HandleMapper<T, To>): Locator<To> {
+    return new MappedLocator(this._clone(), mapper);
+  }
+
   click<ElementType extends Element>(
     this: Locator<ElementType>,
     options?: Readonly<LocatorClickOptions>
diff --git a/packages/puppeteer-core/src/api/locators/MappedLocator.ts b/packages/puppeteer-core/src/api/locators/MappedLocator.ts
index 4d71a9ade8b71..96c515b5cd347 100644
--- a/packages/puppeteer-core/src/api/locators/MappedLocator.ts
+++ b/packages/puppeteer-core/src/api/locators/MappedLocator.ts
@@ -16,7 +16,6 @@
 
 import {Observable, from, mergeMap} from '../../../third_party/rxjs/rxjs.js';
 import {Awaitable, HandleFor} from '../../common/common.js';
-import {JSHandle} from '../JSHandle.js';
 
 import {ActionOptions, DelegatedLocator, Locator} from './locators.js';
 
@@ -25,13 +24,21 @@ import {ActionOptions, DelegatedLocator, Locator} from './locators.js';
  */
 export type Mapper<From, To> = (value: From) => Awaitable<To>;
 
+/**
+ * @internal
+ */
+export type HandleMapper<From, To> = (
+  value: HandleFor<From>,
+  signal?: AbortSignal
+) => Awaitable<HandleFor<To>>;
+
 /**
  * @internal
  */
 export class MappedLocator<From, To> extends DelegatedLocator<From, To> {
-  #mapper: Mapper<From, To>;
+  #mapper: HandleMapper<From, To>;
 
-  constructor(base: Locator<From>, mapper: Mapper<From, To>) {
+  constructor(base: Locator<From>, mapper: HandleMapper<From, To>) {
     super(base);
     this.#mapper = mapper;
   }
@@ -45,7 +52,7 @@ export class MappedLocator<From, To> extends DelegatedLocator<From, To> {
   override _wait(options?: Readonly<ActionOptions>): Observable<HandleFor<To>> {
     return this.delegate._wait(options).pipe(
       mergeMap(handle => {
-        return from((handle as JSHandle<From>).evaluateHandle(this.#mapper));
+        return from(Promise.resolve(this.#mapper(handle, options?.signal)));
       })
     );
   }