feat: allow handling other targets as pages internally

puppeteer · May 12, 2022 · 2e8dad4 · 2e8dad4
1 parent 9ad6fd8
commit 2e8dad4
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 14 deletions.
diff --git a/src/common/Browser.ts b/src/common/Browser.ts
@@ -54,6 +54,13 @@ export type TargetFilterCallback = (
   target: Protocol.Target.TargetInfo
 ) => boolean;
 
+/**
+ * @internal
+ */
+export type IsPageTargetCallback = (
+  target: Protocol.Target.TargetInfo
+) => boolean;
+
 const WEB_PERMISSION_TO_PROTOCOL_PERMISSION = new Map<
   Permission,
   Protocol.Browser.PermissionType
@@ -217,7 +224,8 @@ export class Browser extends EventEmitter {
     defaultViewport?: Viewport | null,
     process?: ChildProcess,
     closeCallback?: BrowserCloseCallback,
-    targetFilterCallback?: TargetFilterCallback
+    targetFilterCallback?: TargetFilterCallback,
+    isPageTargetCallback?: IsPageTargetCallback
   ): Promise<Browser> {
     const browser = new Browser(
       connection,
@@ -226,7 +234,8 @@ export class Browser extends EventEmitter {
       defaultViewport,
       process,
       closeCallback,
-      targetFilterCallback
+      targetFilterCallback,
+      isPageTargetCallback
     );
     await connection.send('Target.setDiscoverTargets', { discover: true });
     return browser;
@@ -237,6 +246,7 @@ export class Browser extends EventEmitter {
   private _connection: Connection;
   private _closeCallback: BrowserCloseCallback;
   private _targetFilterCallback: TargetFilterCallback;
+  private _isPageTargetCallback: IsPageTargetCallback;
   private _defaultContext: BrowserContext;
   private _contexts: Map<string, BrowserContext>;
   private _screenshotTaskQueue: TaskQueue;
@@ -257,7 +267,8 @@ export class Browser extends EventEmitter {
     defaultViewport?: Viewport | null,
     process?: ChildProcess,
     closeCallback?: BrowserCloseCallback,
-    targetFilterCallback?: TargetFilterCallback
+    targetFilterCallback?: TargetFilterCallback,
+    isPageTargetCallback?: IsPageTargetCallback
   ) {
     super();
     this._ignoreHTTPSErrors = ignoreHTTPSErrors;
@@ -267,6 +278,15 @@ export class Browser extends EventEmitter {
     this._connection = connection;
     this._closeCallback = closeCallback || function (): void {};
     this._targetFilterCallback = targetFilterCallback || ((): boolean => true);
+    this._isPageTargetCallback =
+      isPageTargetCallback ||
+      ((target: Protocol.Target.TargetInfo): boolean => {
+        return (
+          target.type === 'page' ||
+          target.type === 'background_page' ||
+          target.type === 'webview'
+        );
+      });
 
     this._defaultContext = new BrowserContext(this._connection, this);
     this._contexts = new Map();
@@ -392,7 +412,8 @@ export class Browser extends EventEmitter {
       () => this._connection.createSession(targetInfo),
       this._ignoreHTTPSErrors,
       this._defaultViewport ?? null,
-      this._screenshotTaskQueue
+      this._screenshotTaskQueue,
+      this._isPageTargetCallback
     );
     assert(
       !this._targets.has(event.targetInfo.targetId),

diff --git a/src/common/Target.ts b/src/common/Target.ts
@@ -17,7 +17,7 @@
 import { Page, PageEmittedEvents } from './Page.js';
 import { WebWorker } from './WebWorker.js';
 import { CDPSession } from './Connection.js';
-import { Browser, BrowserContext } from './Browser.js';
+import { Browser, BrowserContext, IsPageTargetCallback } from './Browser.js';
 import { Viewport } from './PuppeteerViewport.js';
 import { Protocol } from 'devtools-protocol';
 import { TaskQueue } from './TaskQueue.js';
@@ -59,6 +59,10 @@ export class Target {
    * @internal
    */
   _targetId: string;
+  /**
+   * @internal
+   */
+  _isPageTargetCallback: IsPageTargetCallback;
 
   /**
    * @internal
@@ -69,7 +73,8 @@ export class Target {
     sessionFactory: () => Promise<CDPSession>,
     ignoreHTTPSErrors: boolean,
     defaultViewport: Viewport | null,
-    screenshotTaskQueue: TaskQueue
+    screenshotTaskQueue: TaskQueue,
+    isPageTargetCallback: IsPageTargetCallback
   ) {
     this._targetInfo = targetInfo;
     this._browserContext = browserContext;
@@ -78,6 +83,7 @@ export class Target {
     this._ignoreHTTPSErrors = ignoreHTTPSErrors;
     this._defaultViewport = defaultViewport;
     this._screenshotTaskQueue = screenshotTaskQueue;
+    this._isPageTargetCallback = isPageTargetCallback;
     /** @type {?Promise<!Puppeteer.Page>} */
     this._pagePromise = null;
     /** @type {?Promise<!WebWorker>} */
@@ -99,7 +105,8 @@ export class Target {
       (fulfill) => (this._closedCallback = fulfill)
     );
     this._isInitialized =
-      this._targetInfo.type !== 'page' || this._targetInfo.url !== '';
+      !this._isPageTargetCallback(this._targetInfo) ||
+      this._targetInfo.url !== '';
     if (this._isInitialized) this._initializedCallback(true);
   }
 
@@ -114,12 +121,7 @@ export class Target {
    * If the target is not of type `"page"` or `"background_page"`, returns `null`.
    */
   async page(): Promise<Page | null> {
-    if (
-      (this._targetInfo.type === 'page' ||
-        this._targetInfo.type === 'background_page' ||
-        this._targetInfo.type === 'webview') &&
-      !this._pagePromise
-    ) {
+    if (this._isPageTargetCallback(this._targetInfo) && !this._pagePromise) {
       this._pagePromise = this._sessionFactory().then((client) =>
         Page.create(
           client,
@@ -220,7 +222,8 @@ export class Target {
 
     if (
       !this._isInitialized &&
-      (this._targetInfo.type !== 'page' || this._targetInfo.url !== '')
+      (!this._isPageTargetCallback(this._targetInfo) ||
+        this._targetInfo.url !== '')
     ) {
       this._isInitialized = true;
       this._initializedCallback(true);