fix: waitForSelector should work for pseudo classes (#12545)

Co-authored-by: Nikolay Vitkov <34244704+Lightning00Blade@users.noreply.github.com>

packages/puppeteer-core/src/api/ElementHandle.ts

@@ -9,6 +9,7 @@ import type {Protocol} from 'devtools-protocol';
 import type {Frame} from '../api/Frame.js';
 import {getQueryHandlerAndSelector} from '../common/GetQueryHandler.js';
 import {LazyArg} from '../common/LazyArg.js';
+import {PollingOptions} from '../common/QueryHandler.js';
 import type {
   AwaitableIterable,
   ElementFor,
@@ -534,13 +535,12 @@ export abstract class ElementHandle<
     selector: Selector,
     options: WaitForSelectorOptions = {}
   ): Promise<ElementHandle<NodeFor<Selector>> | null> {
-    const {updatedSelector, QueryHandler} =
+    const {updatedSelector, QueryHandler, selectorHasPseudoClasses} =
       getQueryHandlerAndSelector(selector);
-    return (await QueryHandler.waitFor(
+    return (await QueryHandler.waitFor(this, updatedSelector, {
-      this,
+      polling: selectorHasPseudoClasses ? PollingOptions.RAF : undefined,
-      updatedSelector,
+      ...options,
-      options
+    })) as ElementHandle<NodeFor<Selector>> | null;
-    )) as ElementHandle<NodeFor<Selector>> | null;
   }
 
   async #checkVisibility(visibility: boolean): Promise<boolean> {

packages/puppeteer-core/src/api/Frame.ts

@@ -18,6 +18,7 @@ import type {PuppeteerLifeCycleEvent} from '../cdp/LifecycleWatcher.js';
 import {EventEmitter, type EventType} from '../common/EventEmitter.js';
 import {getQueryHandlerAndSelector} from '../common/GetQueryHandler.js';
 import {transposeIterableHandle} from '../common/HandleIterator.js';
+import {PollingOptions} from '../common/QueryHandler.js';
 import type {
   Awaitable,
   EvaluateFunc,
@@ -716,13 +717,12 @@ export abstract class Frame extends EventEmitter<FrameEvents> {
     selector: Selector,
     options: WaitForSelectorOptions = {}
   ): Promise<ElementHandle<NodeFor<Selector>> | null> {
-    const {updatedSelector, QueryHandler} =
+    const {updatedSelector, QueryHandler, selectorHasPseudoClasses} =
       getQueryHandlerAndSelector(selector);
-    return (await QueryHandler.waitFor(
+    return (await QueryHandler.waitFor(this, updatedSelector, {
-      this,
+      polling: selectorHasPseudoClasses ? PollingOptions.RAF : undefined,
-      updatedSelector,
+      ...options,
-      options
+    })) as ElementHandle<NodeFor<Selector>> | null;
-    )) as ElementHandle<NodeFor<Selector>> | null;
   }
 
   /**

packages/puppeteer-core/src/common/GetQueryHandler.ts

@@ -29,6 +29,7 @@ const QUERY_SEPARATORS = ['=', '/'];
  */
 export function getQueryHandlerAndSelector(selector: string): {
   updatedSelector: string;
+  selectorHasPseudoClasses: boolean;
   QueryHandler: typeof QueryHandler;
 } {
   for (const handlerMap of [
@@ -42,20 +43,26 @@ export function getQueryHandlerAndSelector(selector: string): {
         const prefix = `${name}${separator}`;
         if (selector.startsWith(prefix)) {
           selector = selector.slice(prefix.length);
-          return {updatedSelector: selector, QueryHandler};
+          return {
+            updatedSelector: selector,
+            selectorHasPseudoClasses: false,
+            QueryHandler,
+          };
         }
       }
     }
   }
-  const [pSelector, isPureCSS] = parsePSelectors(selector);
+  const [pSelector, isPureCSS, hasPseudoClasses] = parsePSelectors(selector);
   if (isPureCSS) {
     return {
       updatedSelector: selector,
+      selectorHasPseudoClasses: hasPseudoClasses,
       QueryHandler: CSSQueryHandler,
     };
   }
   return {
     updatedSelector: JSON.stringify(pSelector),
+    selectorHasPseudoClasses: hasPseudoClasses,
     QueryHandler: PQueryHandler,
   };
 }

packages/puppeteer-core/src/common/PSelectorParser.ts

@@ -37,11 +37,16 @@ const unquote = (text: string): string => {
  */
 export function parsePSelectors(
   selector: string
-): [selector: ComplexPSelectorList, isPureCSS: boolean] {
+): [
+  selector: ComplexPSelectorList,
+  isPureCSS: boolean,
+  hasPseudoClasses: boolean,
+] {
   let isPureCSS = true;
+  let hasPseudoClasses = false;
   const tokens = tokenize(selector);
   if (tokens.length === 0) {
-    return [[], isPureCSS];
+    return [[], isPureCSS, hasPseudoClasses];
   }
   let compoundSelector: CompoundPSelector = [];
   let complexSelector: ComplexPSelector = [compoundSelector];
@@ -87,6 +92,9 @@ export function parsePSelectors(
           value: unquote(token.argument ?? ''),
         });
         continue;
+      case 'pseudo-class':
+        hasPseudoClasses = true;
+        continue;
       case 'comma':
         if (storage.length) {
           compoundSelector.push(stringify(storage));
@@ -102,5 +110,5 @@ export function parsePSelectors(
   if (storage.length) {
     compoundSelector.push(stringify(storage));
   }
-  return [selectors, isPureCSS];
+  return [selectors, isPureCSS, hasPseudoClasses];
 }

packages/puppeteer-core/src/common/QueryHandler.ts

@@ -34,6 +34,14 @@ export type QuerySelector = (
   PuppeteerUtil: PuppeteerUtil
 ) => Awaitable<Node | null>;
 
+/**
+ * @internal
+ */
+export const enum PollingOptions {
+  RAF = 'raf',
+  MUTATION = 'mutation',
+}
+
 /**
  * @internal
  */
@@ -139,7 +147,9 @@ export class QueryHandler {
   static async waitFor(
     elementOrFrame: ElementHandle<Node> | Frame,
     selector: string,
-    options: WaitForSelectorOptions
+    options: WaitForSelectorOptions & {
+      polling?: PollingOptions;
+    }
   ): Promise<ElementHandle<Node> | null> {
     let frame!: Frame;
     using element = await (async () => {
@@ -152,6 +162,9 @@ export class QueryHandler {
     })();
 
     const {visible = false, hidden = false, timeout, signal} = options;
+    const polling =
+      options.polling ??
+      (visible || hidden ? PollingOptions.RAF : PollingOptions.MUTATION);
 
     try {
       signal?.throwIfAborted();
@@ -169,7 +182,7 @@ export class QueryHandler {
           return PuppeteerUtil.checkVisibility(node, visible);
         },
         {
-          polling: visible || hidden ? 'raf' : 'mutation',
+          polling,
           root: element,
           timeout,
           signal,

test/src/waittask.spec.ts

@@ -407,6 +407,19 @@ describe('waittask specs', function () {
       await watchdog;
     });
 
+    it('should work for selector with a pseudo class', async () => {
+      const {page, server} = await getTestState();
+
+      await page.goto(server.EMPTY_PAGE);
+      const watchdog = page.waitForSelector('input:focus');
+      await expect(
+        Promise.race([watchdog, createTimeout(40)])
+      ).resolves.toBeFalsy();
+      await page.setContent(`<input></input>`);
+      await page.click('input');
+      await watchdog;
+    });
+
     it('Page.waitForSelector is shortcut for main frame', async () => {
       const {page, server} = await getTestState();