From 192edcd269a4c35254b0eebd1ac577c1f35c8e46 Mon Sep 17 00:00:00 2001
From: Larry1123 <larry1123@larry1123.net>
Date: Fri, 14 Sep 2018 13:34:19 -0400
Subject: [PATCH] Extend screenshot encoding support (#28713)

Overloaded the screenshot methods to define the return type based on what kind of encoding is used.
Edited the test and made sure normal ScreenshotOptions worked as normal still, the two new extended interfaces can be used for objects to pass in also but did not add test for that directly, should be handled right with the current test cases.
---
 types/puppeteer/index.d.ts         | 14 +++++++++++++-
 types/puppeteer/puppeteer-tests.ts |  8 ++++++--
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/types/puppeteer/index.d.ts b/types/puppeteer/index.d.ts
index e66ee97dfa..eab93cf2eb 100644
--- a/types/puppeteer/index.d.ts
+++ b/types/puppeteer/index.d.ts
@@ -592,6 +592,14 @@ export interface ScreenshotOptions {
   encoding?: "base64" | "binary";
 }
 
+export interface BinaryScreenShotOptions extends ScreenshotOptions {
+    encoding?: "binary";
+}
+
+export interface Base64ScreenShotOptions extends ScreenshotOptions {
+    encoding: "base64";
+}
+
 /** Options for `addStyleTag` */
 export interface StyleTagOptions {
   /** Url of the <link> tag. */
@@ -747,7 +755,9 @@ export interface ElementHandle<E extends Element = Element> extends JSHandle, Ev
    * If the element is detached from DOM, the method throws an error.
    * @param options Same options as in page.screenshot.
    */
-  screenshot(options?: ScreenshotOptions): Promise<Buffer>;
+  screenshot(options?: Base64ScreenShotOptions): Promise<string>;
+  screenshot(options?: BinaryScreenShotOptions): Promise<Buffer>;
+  screenshot(options?: ScreenshotOptions): Promise<string | Buffer>;
   /**
    * This method scrolls element into view if needed, and then uses touchscreen.tap to tap in the center of the element.
    * If the element is detached from DOM, the method throws an error.
@@ -1346,6 +1356,8 @@ export interface Page extends EventEmitter, FrameBase {
    * Captures a screenshot of the page.
    * @param options The screenshot options.
    */
+  screenshot(options?: Base64ScreenShotOptions): Promise<string>;
+  screenshot(options?: BinaryScreenShotOptions): Promise<Buffer>;
   screenshot(options?: ScreenshotOptions): Promise<string | Buffer>;
 
   /**
diff --git a/types/puppeteer/puppeteer-tests.ts b/types/puppeteer/puppeteer-tests.ts
index 1e3d35737b..fb0c283d4b 100644
--- a/types/puppeteer/puppeteer-tests.ts
+++ b/types/puppeteer/puppeteer-tests.ts
@@ -403,8 +403,12 @@ puppeteer.launch().then(async browser => {
   const browser = await puppeteer.launch();
   const page = await browser.newPage();
   await page.goto("https://example.com");
-  const base64string: string = (await page.screenshot({ encoding: "base64" })) as string;
-  const buffer: Buffer = (await page.screenshot({ encoding: "binary" })) as Buffer;
+  const base64string: string = await page.screenshot({ encoding: "base64" });
+  const buffer: Buffer = await page.screenshot({ encoding: "binary" });
+  const screenshotOptions: puppeteer.ScreenshotOptions = {
+    fullPage: true,
+  };
+  const stringOrBuffer: string | Buffer = await page.screenshot(screenshotOptions);
 
   browser.close();
 })();