From 23c0ba07276c301e2bec6c5c805f618d0c17fccc Mon Sep 17 00:00:00 2001
From: Andrey Lushnikov <aslushnikov@gmail.com>
Date: Wed, 11 Oct 2017 14:41:20 -0700
Subject: [PATCH] feat(Page): introduce Page.queryObjects (#1005)

This patch introduces `Page.queryObjects` and
`ExecutionContext.queryObjects` methods to query JavaScript heap
for objects with a certain prototype.

Fixes #304.
---
 docs/api.md             | 42 +++++++++++++++++++++++++++++++++++++++++
 lib/ExecutionContext.js | 13 +++++++++++++
 lib/Page.js             |  8 ++++++++
 test/test.js            | 26 +++++++++++++++++++++++++
 4 files changed, 89 insertions(+)

diff --git a/docs/api.md b/docs/api.md
index b321df43..79385097 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -61,6 +61,7 @@
     + [page.mainFrame()](#pagemainframe)
     + [page.mouse](#pagemouse)
     + [page.pdf(options)](#pagepdfoptions)
+    + [page.queryObjects(prototypeHandle)](#pagequeryobjectsprototypehandle)
     + [page.reload(options)](#pagereloadoptions)
     + [page.screenshot([options])](#pagescreenshotoptions)
     + [page.select(selector, ...values)](#pageselectselector-values)
@@ -130,6 +131,7 @@
   * [class: ExecutionContext](#class-executioncontext)
     + [executionContext.evaluate(pageFunction, ...args)](#executioncontextevaluatepagefunction-args)
     + [executionContext.evaluateHandle(pageFunction, ...args)](#executioncontextevaluatehandlepagefunction-args)
+    + [executionContext.queryObjects(prototypeHandle)](#executioncontextqueryobjectsprototypehandle)
   * [class: JSHandle](#class-jshandle)
     + [jsHandle.asElement()](#jshandleaselement)
     + [jsHandle.dispose()](#jshandledispose)
@@ -811,6 +813,27 @@ The `format` options are:
 - `A5`: 5.83in x 8.27in
 - `A6`: 4.13in x 5.83in
 
+#### page.queryObjects(prototypeHandle)
+- `prototypeHandle` <[JSHandle]> A handle to the object prototype.
+- returns: <[JSHandle]> A handle to an array of objects with this prototype
+
+The method iterates javascript heap and finds all the objects with the given prototype.
+
+```js
+// Create a Map object 
+await page.evaluate(() => window.map = new Map());
+// Get a handle to the Map object prototype
+const mapPrototype = await page.evaluateHandle(() => Map.prototype);
+// Query all map instances into an array
+const mapInstances = await page.queryObjects(mapPrototype);
+// Count amount of map objects in heap
+const count = await page.evaluate(maps => maps.length, mapInstances);
+await mapInstances.dispose();
+await mapPrototype.dispose();
+```
+
+Shortcut for [page.mainFrame().executionContext().queryObjects(prototypeHandle)](#executioncontextqueryobjectsprototypehandle).
+
 #### page.reload(options)
 - `options` <[Object]> Navigation parameters which might have the following properties:
   - `timeout` <[number]> Maximum navigation time in milliseconds, defaults to 30 seconds, pass `0` to disable timeout.
@@ -1513,6 +1536,25 @@ await aHandle.dispose();
 await resultHandle.dispose();
 ```
 
+#### executionContext.queryObjects(prototypeHandle)
+- `prototypeHandle` <[JSHandle]> A handle to the object prototype.
+- returns: <[JSHandle]> A handle to an array of objects with this prototype
+
+The method iterates javascript heap and finds all the objects with the given prototype.
+
+```js
+// Create a Map object 
+await page.evaluate(() => window.map = new Map());
+// Get a handle to the Map object prototype
+const mapPrototype = await page.evaluateHandle(() => Map.prototype);
+// Query all map instances into an array
+const mapInstances = await page.queryObjects(mapPrototype);
+// Count amount of map objects in heap
+const count = await page.evaluate(maps => maps.length, mapInstances);
+await mapInstances.dispose();
+await mapPrototype.dispose();
+```
+
 ### class: JSHandle
 
 JSHandle represents an in-page javascript object. JSHandles can be created with the [page.evaluateHandle](#pageevaluatehandlepagefunction-args) method.
diff --git a/lib/ExecutionContext.js b/lib/ExecutionContext.js
index 7fc3f22c..325cd801 100644
--- a/lib/ExecutionContext.js
+++ b/lib/ExecutionContext.js
@@ -95,6 +95,19 @@ class ExecutionContext {
       return { value: arg };
     }
   }
+
+  /**
+   * @param {!JSHandle} prototypeHandle
+   * @return {!Promise<!JSHandle>}
+   */
+  async queryObjects(prototypeHandle) {
+    console.assert(!prototypeHandle._disposed, 'Prototype JSHandle is disposed!');
+    console.assert(prototypeHandle._remoteObject.objectId, 'Prototype JSHandle must not be referencing primitive value');
+    const response = await this._client.send('Runtime.queryObjects', {
+      prototypeObjectId: prototypeHandle._remoteObject.objectId
+    });
+    return this._objectHandleFactory(response.objects);
+  }
 }
 
 class JSHandle {
diff --git a/lib/Page.js b/lib/Page.js
index 7cccaa4b..ff72d2ad 100644
--- a/lib/Page.js
+++ b/lib/Page.js
@@ -178,6 +178,14 @@ class Page extends EventEmitter {
     return this.mainFrame().executionContext().evaluateHandle(pageFunction, ...args);
   }
 
+  /**
+   * @param {!Puppeteer.JSHandle} prototypeHandle
+   * @return {!Promise<!Puppeteer.JSHandle>}
+   */
+  async queryObjects(prototypeHandle) {
+    return this.mainFrame().executionContext().queryObjects(prototypeHandle);
+  }
+
   /**
    * @param {string} selector
    * @param {function()|string} pageFunction
diff --git a/test/test.js b/test/test.js
index 88d4f67a..a7e1e294 100644
--- a/test/test.js
+++ b/test/test.js
@@ -335,6 +335,32 @@ describe('Page', function() {
     }));
   });
 
+  describe('ExecutionContext.queryObjects', function() {
+    it('should work', SX(async function() {
+      // Instantiate an object
+      await page.evaluate(() => window.set = new Set(['hello', 'world']));
+      const prototypeHandle = await page.evaluateHandle(() => Set.prototype);
+      const objectsHandle = await page.queryObjects(prototypeHandle);
+      const count = await page.evaluate(objects => objects.length, objectsHandle);
+      expect(count).toBe(1);
+      const values = await page.evaluate(objects => Array.from(objects[0].values()), objectsHandle);
+      expect(values).toEqual(['hello', 'world']);
+    }));
+    it('should fail for disposed handles', SX(async function() {
+      const prototypeHandle = await page.evaluateHandle(() => HTMLBodyElement.prototype);
+      await prototypeHandle.dispose();
+      let error = null;
+      await page.queryObjects(prototypeHandle).catch(e => error = e);
+      expect(error.message).toBe('Prototype JSHandle is disposed!');
+    }));
+    it('should fail primitive values as prototypes', SX(async function() {
+      const prototypeHandle = await page.evaluateHandle(() => 42);
+      let error = null;
+      await page.queryObjects(prototypeHandle).catch(e => error = e);
+      expect(error.message).toBe('Prototype JSHandle must not be referencing primitive value');
+    }));
+  });
+
   describe('JSHandle.getProperty', function() {
     it('should work', SX(async function() {
       const aHandle = await page.evaluateHandle(() => ({