chore: update deep implementation for P selectors (#9908)

packages/puppeteer-core/src/injected/PQuerySelector.ts

@@ -29,11 +29,19 @@ import {
   PPseudoSelector,
 } from './PSelectorParser.js';
 import {textQuerySelectorAll} from './TextQuerySelector.js';
-import {deepChildren, deepDescendents} from './util.js';
+import {pierce, pierceAll} from './util.js';
 import {xpathQuerySelectorAll} from './XPathQuerySelector.js';
 
 const IDENT_TOKEN_START = /[-\w\P{ASCII}*]/;
 
+interface QueryableNode extends Node {
+  querySelectorAll: typeof Document.prototype.querySelectorAll;
+}
+
+const isQueryableNode = (node: Node): node is QueryableNode => {
+  return 'querySelectorAll' in node;
+};
+
 class SelectorError extends Error {
   constructor(selector: string, message: string) {
     super(`${selector} is not a valid selector: ${message}`);
@@ -75,33 +83,44 @@ class PQueryEngine {
       const selector = this.#selector;
       const input = this.#input;
       if (typeof selector === 'string') {
-        this.elements = AsyncIterableUtil.flatMap(
-          this.elements,
-          async function* (element) {
-            if (!selector[0]) {
-              return;
-            }
-            // The regular expression tests if the selector is a type/universal
-            // selector. Any other case means we want to apply the selector onto
-            // the element itself (e.g. `element.class`, `element>div`,
-            // `element:hover`, etc.).
-            if (IDENT_TOKEN_START.test(selector[0]) || !element.parentElement) {
-              yield* (element as Element).querySelectorAll(selector);
-              return;
-            }
-
-            let index = 0;
-            for (const child of element.parentElement.children) {
-              ++index;
-              if (child === element) {
-                break;
+        // The regular expression tests if the selector is a type/universal
+        // selector. Any other case means we want to apply the selector onto
+        // the element itself (e.g. `element.class`, `element>div`,
+        // `element:hover`, etc.).
+        if (selector[0] && IDENT_TOKEN_START.test(selector[0])) {
+          this.elements = AsyncIterableUtil.flatMap(
+            this.elements,
+            async function* (element) {
+              if (isQueryableNode(element)) {
+                yield* element.querySelectorAll(selector);
               }
             }
-            yield* element.parentElement.querySelectorAll(
-              `:scope>:nth-child(${index})${selector}`
-            );
-          }
-        );
+          );
+        } else {
+          this.elements = AsyncIterableUtil.flatMap(
+            this.elements,
+            async function* (element) {
+              if (!element.parentElement) {
+                if (!isQueryableNode(element)) {
+                  return;
+                }
+                yield* element.querySelectorAll(selector);
+                return;
+              }
+
+              let index = 0;
+              for (const child of element.parentElement.children) {
+                ++index;
+                if (child === element) {
+                  break;
+                }
+              }
+              yield* element.parentElement.querySelectorAll(
+                `:scope>:nth-child(${index})${selector}`
+              );
+            }
+          );
+        }
       } else {
         this.elements = AsyncIterableUtil.flatMap(
           this.elements,
@@ -144,22 +163,12 @@ class PQueryEngine {
     const selector = this.#complexSelector.shift();
     switch (selector) {
       case PCombinator.Child: {
-        this.elements = AsyncIterableUtil.flatMap(
-          this.elements,
-          function* (element) {
-            yield* deepChildren(element);
-          }
-        );
+        this.elements = AsyncIterableUtil.flatMap(this.elements, pierce);
         this.#next();
         break;
       }
       case PCombinator.Descendent: {
-        this.elements = AsyncIterableUtil.flatMap(
-          this.elements,
-          function* (element) {
-            yield* deepDescendents(element);
-          }
-        );
+        this.elements = AsyncIterableUtil.flatMap(this.elements, pierceAll);
         this.#next();
         break;
       }
@@ -206,12 +215,12 @@ const compareDepths = (a: number[], b: number[]): -1 | 0 | 1 => {
   if (a.length + b.length === 0) {
     return 0;
   }
-  const [i = Infinity, ...otherA] = a;
-  const [j = Infinity, ...otherB] = b;
+  const [i = -1, ...otherA] = a;
+  const [j = -1, ...otherB] = b;
   if (i === j) {
     return compareDepths(otherA, otherB);
   }
-  return i < j ? 1 : -1;
+  return i < j ? -1 : 1;
 };
 
 const domSort = async function* (elements: AwaitableIterable<Node>) {
@@ -232,10 +241,6 @@ const domSort = async function* (elements: AwaitableIterable<Node>) {
     });
 };
 
-type QueryableNode = {
-  querySelectorAll: typeof Document.prototype.querySelectorAll;
-};
-
 /**
  * Queries the given node for all nodes matching the given text selector.
  *

packages/puppeteer-core/src/injected/util.ts

@@ -30,41 +30,37 @@ function isBoundingBoxEmpty(element: Element): boolean {
   return rect.width === 0 || rect.height === 0;
 }
 
+const hasShadowRoot = (node: Node): node is Node & {shadowRoot: ShadowRoot} => {
+  return 'shadowRoot' in node && node.shadowRoot instanceof ShadowRoot;
+};
+
 /**
  * @internal
  */
-export function* deepChildren(
-  root: Node
-): IterableIterator<Element | ShadowRoot> {
-  const walker = document.createTreeWalker(root, NodeFilter.SHOW_ELEMENT);
-  let node = walker.nextNode() as Element | null;
-  for (; node; node = walker.nextNode() as Element | null) {
-    yield node.shadowRoot ?? node;
+export function* pierce(root: Node): IterableIterator<Node | ShadowRoot> {
+  if (hasShadowRoot(root)) {
+    yield root.shadowRoot;
+  } else {
+    yield root;
   }
 }
 
 /**
  * @internal
  */
-export function* deepDescendents(
-  root: Node
-): IterableIterator<Element | ShadowRoot> {
+export function* pierceAll(root: Node): IterableIterator<Node | ShadowRoot> {
+  yield* pierce(root);
   const walkers = [document.createTreeWalker(root, NodeFilter.SHOW_ELEMENT)];
-  let walker: TreeWalker | undefined;
-  while ((walker = walkers.shift())) {
-    for (
-      let node = walker.nextNode() as Element | null;
-      node;
-      node = walker.nextNode() as Element | null
-    ) {
+  for (const walker of walkers) {
+    let node: Element | null;
+    while ((node = walker.nextNode() as Element | null)) {
       if (!node.shadowRoot) {
-        yield node;
         continue;
       }
+      yield node.shadowRoot;
       walkers.push(
         document.createTreeWalker(node.shadowRoot, NodeFilter.SHOW_ELEMENT)
       );
-      yield node.shadowRoot;
     }
   }
 }

test/assets/p-selectors.html

@@ -0,0 +1,13 @@
+<div id="a">hello <button id="b">world</button>
+    <span id="f"></span>
+    <div id="c">
+        <template shadowrootmode="open">
+            shadow dom
+            <div id="d">
+                <template shadowrootmode="open">
+                    <a id="e">deep text</a>
+                </template>
+            </div>
+        </template>
+    </div>
+</div>

test/src/queryhandler.spec.ts

@@ -358,10 +358,8 @@ describe('Query handler tests', function () {
 
   describe('P selectors', () => {
     beforeEach(async () => {
-      const {page} = getTestState();
-      await page.setContent(
-        '<div>hello <button>world<span></span></button></div>'
-      );
+      const {page, server} = getTestState();
+      await page.goto(`${server.PREFIX}/p-selectors.html`);
       Puppeteer.clearCustomQueryHandlers();
     });
 
@@ -371,7 +369,7 @@ describe('Query handler tests', function () {
       assert(element, 'Could not find element');
       expect(
         await element.evaluate(element => {
-          return element.tagName === 'BUTTON';
+          return element.id === 'b';
         })
       ).toBeTruthy();
 
@@ -386,13 +384,35 @@ describe('Query handler tests', function () {
       }
     });
 
+    it('should work with deep combinators', async () => {
+      const {page} = getTestState();
+      {
+        const element = await page.$('div >>>> div');
+        assert(element, 'Could not find element');
+        expect(
+          await element.evaluate(element => {
+            return element.id === 'c';
+          })
+        ).toBeTruthy();
+      }
+      {
+        const elements = await page.$$('div >>> div');
+        assert(elements[1], 'Could not find element');
+        expect(
+          await elements[1]?.evaluate(element => {
+            return element.id === 'd';
+          })
+        ).toBeTruthy();
+      }
+    });
+
     it('should work with text selectors', async () => {
       const {page} = getTestState();
       const element = await page.$('div ::-p-text(world)');
       assert(element, 'Could not find element');
       expect(
         await element.evaluate(element => {
-          return element.tagName === 'BUTTON';
+          return element.id === 'b';
         })
       ).toBeTruthy();
     });
@@ -403,7 +423,7 @@ describe('Query handler tests', function () {
       assert(element, 'Could not find element');
       expect(
         await element.evaluate(element => {
-          return element.tagName === 'BUTTON';
+          return element.id === 'b';
         })
       ).toBeTruthy();
     });
@@ -414,7 +434,7 @@ describe('Query handler tests', function () {
       assert(element, 'Could not find element');
       expect(
         await element.evaluate(element => {
-          return element.tagName === 'BUTTON';
+          return element.id === 'b';
         })
       ).toBeTruthy();
     });
@@ -431,7 +451,7 @@ describe('Query handler tests', function () {
       assert(element, 'Could not find element');
       expect(
         await element.evaluate(element => {
-          return element.tagName === 'DIV';
+          return element.id === 'a';
         })
       ).toBeTruthy();
     });
@@ -453,7 +473,7 @@ describe('Query handler tests', function () {
         assert(element, 'Could not find element');
         expect(
           await element.evaluate(element => {
-            return element.tagName === 'DIV';
+            return element.id === 'a';
           })
         ).toBeTruthy();
       }
@@ -462,7 +482,7 @@ describe('Query handler tests', function () {
         assert(element, 'Could not find element');
         expect(
           await element.evaluate(element => {
-            return element.tagName === 'DIV';
+            return element.id === 'a';
           })
         ).toBeTruthy();
       }
@@ -471,7 +491,7 @@ describe('Query handler tests', function () {
         assert(element, 'Could not find element');
         expect(
           await element.evaluate(element => {
-            return element.tagName === 'DIV';
+            return element.id === 'a';
           })
         ).toBeTruthy();
       }
@@ -480,7 +500,7 @@ describe('Query handler tests', function () {
         assert(element, 'Could not find element');
         expect(
           await element.evaluate(element => {
-            return element.tagName === 'BUTTON';
+            return element.id === 'b';
           })
         ).toBeTruthy();
       }
@@ -504,7 +524,7 @@ describe('Query handler tests', function () {
     it('should work with selector lists', async () => {
       const {page} = getTestState();
       const elements = await page.$$('div, ::-p-text(world)');
-      expect(elements.length).toStrictEqual(2);
+      expect(elements.length).toStrictEqual(3);
     });
 
     const permute = <T>(inputs: T[]): T[][] => {
@@ -528,11 +548,6 @@ describe('Query handler tests', function () {
     it('should match querySelector* ordering', async () => {
       const {page} = getTestState();
       for (const list of permute(['div', 'button', 'span'])) {
-        const expected = await page.evaluate(selector => {
-          return [...document.querySelectorAll(selector)].map(element => {
-            return element.tagName;
-          });
-        }, list.join(','));
         const elements = await page.$$(
           list
             .map(selector => {
@@ -543,11 +558,11 @@ describe('Query handler tests', function () {
         const actual = await Promise.all(
           elements.map(element => {
             return element.evaluate(element => {
-              return element.tagName;
+              return element.id;
             });
           })
         );
-        expect(actual.join()).toStrictEqual(expected.join());
+        expect(actual.join()).toStrictEqual('a,b,f,c');
       }
     });