From 151d512ae2d652d8bc91db1ebcb635fd24a45d1a Mon Sep 17 00:00:00 2001
From: Andrey Lushnikov <aslushnikov@gmail.com>
Date: Tue, 22 Aug 2017 22:56:55 -0700
Subject: [PATCH] Implement page.$$ method (#463)

This patch implements page.$$ which runs document.querySelectorAll
in page and returns results as an array of ElementHandle instances.

Fixes #384.
---
 docs/api.md         | 19 +++++++++++++++++--
 lib/FrameManager.js | 25 ++++++++++++++++++++++++-
 lib/Page.js         |  8 ++++++++
 test/test.js        | 14 ++++++++++++++
 4 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/docs/api.md b/docs/api.md
index 7f45b43b9f8..cd666c89178 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -27,6 +27,7 @@
     + [event: 'requestfinished'](#event-requestfinished)
     + [event: 'response'](#event-response)
     + [page.$(selector)](#pageselector)
+    + [page.$$(selector)](#pageselector)
     + [page.addScriptTag(url)](#pageaddscripttagurl)
     + [page.click(selector[, options])](#pageclickselector-options)
     + [page.close()](#pageclose)
@@ -85,6 +86,7 @@
     + [dialog.type](#dialogtype)
   * [class: Frame](#class-frame)
     + [frame.$(selector)](#frameselector)
+    + [frame.$$(selector)](#frameselector)
     + [frame.addScriptTag(url)](#frameaddscripttagurl)
     + [frame.childFrames()](#framechildframes)
     + [frame.evaluate(pageFunction, ...args)](#frameevaluatepagefunction-args)
@@ -297,6 +299,14 @@ The method runs `document.querySelector` within the page. If no element matches
 
 Shortcut for [page.mainFrame().$(selector)](#frameselector).
 
+#### page.$$(selector)
+- `selector` <[string]> Selector to query page for
+- returns: <[Promise]<[Array]<[ElementHandle]>>>
+
+The method runs `document.querySelectorAll` within the page. If no elements match the selector, the return value resolve to `[]`.
+
+Shortcut for [page.mainFrame().$$(selector)](#frameselector-1).
+
 #### page.addScriptTag(url)
 - `url` <[string]> Url of the `<script>` tag
 - returns: <[Promise]> which resolves when the script's onload fires.
@@ -976,10 +986,15 @@ puppeteer.launch().then(async browser => {
 
 #### frame.$(selector)
 - `selector` <[string]> Selector to query page for
-- returns: <[Promise]<[ElementHandle]>> Promise which resolves to ElementHandle pointing to the page element.
+- returns: <[Promise]<[ElementHandle]>> Promise which resolves to ElementHandle pointing to the frame element.
 
-The method queries page for the selector. If there's no such element within the page, the method will resolve to `null`.
+The method queries frame for the selector. If there's no such element within the frame, the method will resolve to `null`.
 
+#### frame.$$(selector)
+- `selector` <[string]> Selector to query page for
+- returns: <[Promise]<[Array]<[ElementHandle]>>> Promise which resolves to ElementHandles pointing to the frame elements.
+
+The method runs `document.querySelectorAll` within the frame. If no elements match the selector, the return value resolve to `[]`.
 
 #### frame.addScriptTag(url)
 - `url` <[string]> Url of a script to be added
diff --git a/lib/FrameManager.js b/lib/FrameManager.js
index adbffe1fba8..2e25f8c4634 100644
--- a/lib/FrameManager.js
+++ b/lib/FrameManager.js
@@ -191,10 +191,33 @@ class Frame {
     const remoteObject = await this._rawEvaluate(selector => document.querySelector(selector), selector);
     if (remoteObject.subtype === 'node')
       return new ElementHandle(this._client, remoteObject, this._mouse);
-    helper.releaseObject(this._client, remoteObject);
+    await helper.releaseObject(this._client, remoteObject);
     return null;
   }
 
+  /**
+   * @param {string} selector
+   * @return {!Promise<!Array<!ElementHandle>>}
+   */
+  async $$(selector) {
+    const remoteObject = await this._rawEvaluate(selector => Array.from(document.querySelectorAll(selector)), selector);
+    const response = await this._client.send('Runtime.getProperties', {
+      objectId: remoteObject.objectId,
+      ownProperties: true
+    });
+    const properties = response.result;
+    const result = [];
+    const releasePromises = [helper.releaseObject(this._client, remoteObject)];
+    for (const property of properties) {
+      if (property.enumerable && property.value.subtype === 'node')
+        result.push(new ElementHandle(this._client, property.value, this._mouse));
+      else
+        releasePromises.push(helper.releaseObject(this._client, property.value));
+    }
+    await Promise.all(releasePromises);
+    return result;
+  }
+
   /**
    * @param {function()|string} pageFunction
    * @param {!Array<*>} args
diff --git a/lib/Page.js b/lib/Page.js
index c0c26cc84df..35fb92781e8 100644
--- a/lib/Page.js
+++ b/lib/Page.js
@@ -146,6 +146,14 @@ class Page extends EventEmitter {
     return this.mainFrame().$(selector);
   }
 
+  /**
+   * @param {string} selector
+   * @return {!Promise<!Array<!ElementHandle>>}
+   */
+  async $$(selector) {
+    return this.mainFrame().$$(selector);
+  }
+
   /**
    * @param {string} url
    * @return {!Promise}
diff --git a/test/test.js b/test/test.js
index 694915c7643..261c632e5c7 100644
--- a/test/test.js
+++ b/test/test.js
@@ -1074,6 +1074,20 @@ describe('Page', function() {
       expect(element).toBe(null);
     }));
   });
+  describe('Page.$$', function() {
+    it('should query existing elements', SX(async function() {
+      await page.setContent('<div>A</div><br/><div>B</div>');
+      const elements = await page.$$('div');
+      expect(elements.length).toBe(2);
+      const promises = elements.map(element => element.evaluate(e => e.textContent));
+      expect(await Promise.all(promises)).toEqual(['A', 'B']);
+    }));
+    it('should return ampty array if nothing is found', SX(async function() {
+      await page.goto(EMPTY_PAGE);
+      const elements = await page.$$('div');
+      expect(elements.length).toBe(0);
+    }));
+  });
 
   describe('ElementHandle.evaluate', function() {
     it('should work', SX(async function() {