✨ Add beforeScreenshot parameter (#65)

frinyvonnick · Sep 21, 2020 · 5fcb054 · 5fcb054
1 parent be23aa7
commit 5fcb054
Show file tree

Hide file tree

Showing 5 changed files with 81 additions and 49 deletions.
diff --git a/README.md b/README.md
@@ -73,6 +73,7 @@ List of all available options:
 | content                 | If provided html property is considered an handlebars template and use content value to fill it | object or Array            | optional    |
 | waitUntil               | Define when to consider markup succeded. [Learn more](https://github.com/puppeteer/puppeteer/blob/8370ec88ae94fa59d9e9dc0c154e48527d48c9fe/docs/api.md#pagesetcontenthtml-options).                                                        | string or Array<string> (default: networkidle0)    | optional    |
 | puppeteerArgs           | The puppeteerArgs property let you pass down custom configuration to puppeteer. [Learn more](https://github.com/puppeteer/puppeteer/blob/8370ec88ae94fa59d9e9dc0c154e48527d48c9fe/docs/api.md#puppeteerlaunchoptions).                  | object                     | optional    |
+| beforeScreenshot | An async function that will execute just before screenshot is taken. Gives access to puppeteer page element. | Function | optional |
 | transparent             | The transparent property lets you generate images with transparent background (for png type).    | boolean                    | optional    |
 | encoding             | The encoding property of the image. Options are `binary` (default) or `base64`.    | string                    | optional    |
 

diff --git a/src/screenshot.js b/src/screenshot.js
@@ -1,13 +1,14 @@
 const handlebars = require('handlebars')
 
 module.exports = {
-  makeScreenshot: async function(page, {
+  makeScreenshot: async function (page, {
     output,
     type,
     quality,
     encoding,
     content,
     html,
+    beforeScreenshot,
     transparent = false,
     waitUntil = 'networkidle0',
   }) {
@@ -22,6 +23,9 @@ module.exports = {
     }
     await page.setContent(html, { waitUntil })
     const element = await page.$('body')
+    if (beforeScreenshot && typeof beforeScreenshot === "function") {
+      await beforeScreenshot(page);
+    }
     const buffer = await element.screenshot({ path: output, type, omitBackground: transparent, encoding, ...screeshotArgs })
 
     return buffer

diff --git a/src/screenshot.spec.js b/src/screenshot.spec.js
@@ -1,3 +1,4 @@
+const { describe } = require('jest-circus')
 const puppeteer = require('puppeteer')
 const { makeScreenshot } = require('./screenshot')
 
@@ -18,8 +19,8 @@ describe('quality', () => {
       quality: 300,
       html: '<html><body>Hello world!</body></html>',
     })
-    
-    expect(screenshot).toHaveBeenCalledWith(expect.objectContaining({'encoding': undefined, 'omitBackground': false, 'path': undefined, 'type': 'png'}))
+
+    expect(screenshot).toHaveBeenCalledWith(expect.objectContaining({ 'encoding': undefined, 'omitBackground': false, 'path': undefined, 'type': 'png' }))
   })
 
   it('should set quality option for jpg images', async () => {
@@ -28,7 +29,7 @@ describe('quality', () => {
       quality: 30,
       html: '<html><body>Hello world!</body></html>',
     })
-    
+
     expect(screenshot).toHaveBeenCalledWith(expect.objectContaining({ quality: 30 }))
   })
 
@@ -37,11 +38,33 @@ describe('quality', () => {
       type: 'jpeg',
       html: '<html><body>Hello world!</body></html>',
     })
-    
+
     expect(screenshot).toHaveBeenCalledWith(expect.objectContaining({ quality: 80 }))
   })
 })
 
+describe('beforeScreenshot', () => {
+  let screenshot
+  let puppeteer
+  let page
+
+  beforeEach(() => {
+    puppeteer = require('puppeteer')
+    page = puppeteer.launch().newPage()
+    screenshot = page.$().screenshot
+  })
+
+  it('beforeScreenshot is called with page', async () => {
+    const beforeScreenshot = jest.fn();
+    await makeScreenshot(page, {
+      beforeScreenshot,
+      html: '<html><body>Hello world!</body></html>',
+    })
+
+    expect(beforeScreenshot).toHaveBeenCalledWith(page);
+  })
+})
+
 jest.mock('puppeteer', () => {
   const screenshot = jest.fn()
   return {

diff --git a/types/options.d.ts b/types/options.d.ts
@@ -4,46 +4,50 @@ import type { LaunchOptions, LoadEvent } from 'puppeteer';
  * Available options to pass to the library
  */
 export interface NodeHtmlToImageOptions {
-    /**
-     * The html used to generate image content
-     */
-    html: string;
-    /**
-     * The ouput path for generated image
-     * If none is specified
-     */
-    output?: string;
-    /**
-     * The type of the generated image
-     * @default png
-     */
-    type?: 'jpeg' | 'png';
-    /**
-     * The quality of the generated image (only applicable to jpg)
-     * @default 80
-     */
-    quality?: number;
-    /**
-     * If provided html property is considered an handlebars template and use content value to fill it
-     */
-    content?: Array<unknown> | Record<string, any>;
-    /**
-     * Define when to consider markup succeded.
-     * {@link https://github.com/puppeteer/puppeteer/blob/8370ec88ae94fa59d9e9dc0c154e48527d48c9fe/docs/api.md#pagesetcontenthtml-options Learn more}
-     */
-    waitUntil?: LoadEvent;
-    /**
-     * The puppeteerArgs property let you pass down custom configuration to puppeteer.
-     * {@link https://github.com/puppeteer/puppeteer/blob/8370ec88ae94fa59d9e9dc0c154e48527d48c9fe/docs/api.md#puppeteerlaunchoptions Learn more}
-     */
-    puppeteerArgs?: LaunchOptions;
-    /**
-     * The transparent property lets you generate images with transparent background (for png type).
-     */
-    transparent?: boolean;
-    /**
-     * The encoding property of the image. Options are `binary` (default) or `base64`.
-     * @default binary
-     */
-    encoding?: 'binary' | 'base64';
-}
+  /**
+   * The html used to generate image content
+   */
+  html: string;
+  /**
+   * The ouput path for generated image
+   * If none is specified
+   */
+  output?: string;
+  /**
+   * The type of the generated image
+   * @default png
+   */
+  type?: 'jpeg' | 'png';
+  /**
+   * The quality of the generated image (only applicable to jpg)
+   * @default 80
+   */
+  quality?: number;
+  /**
+   * If provided html property is considered an handlebars template and use content value to fill it
+   */
+  content?: Array<unknown> | Record<string, any>;
+  /**
+   * Define when to consider markup succeded.
+   * {@link https://github.com/puppeteer/puppeteer/blob/8370ec88ae94fa59d9e9dc0c154e48527d48c9fe/docs/api.md#pagesetcontenthtml-options Learn more}
+   */
+  waitUntil?: LoadEvent;
+  /**
+   * The puppeteerArgs property let you pass down custom configuration to puppeteer.
+   * {@link https://github.com/puppeteer/puppeteer/blob/8370ec88ae94fa59d9e9dc0c154e48527d48c9fe/docs/api.md#puppeteerlaunchoptions Learn more}
+   */
+  puppeteerArgs?: LaunchOptions;
+  /**
+   * Function executes before screenshot is taken and provides access to puppeteer page element
+   */
+  beforeScreenshot?: Function;
+  /**
+   * The transparent property lets you generate images with transparent background (for png type).
+   */
+  transparent?: boolean;
+  /**
+   * The encoding property of the image. Options are `binary` (default) or `base64`.
+   * @default binary
+   */
+  encoding?: 'binary' | 'base64';
+}
diff --git a/types/screenshot.d.ts b/types/screenshot.d.ts
@@ -2,5 +2,5 @@
 import type { Page } from 'puppeteer';
 import type { NodeHtmlToImageOptions } from './options';
 
-declare function makeScreenshot(page: Page, { output, type, quality, encoding, content, html, transparent, waitUntil }: NodeHtmlToImageOptions): Promise<string | Buffer>;
+declare function makeScreenshot(page: Page, { output, type, quality, encoding, content, html, beforeScreenshot, transparent, waitUntil }: NodeHtmlToImageOptions): Promise<string | Buffer>;
 export default makeScreenshot;