Add selector option (#114)

Co-authored-by: derek <derek@cemware.com> Co-authored-by: Yvonnick Frin <frin.yvonnick@gmail.com>
frinyvonnick · Jun 26, 2021 · 701b514 · 701b514
1 parent a9db417
commit 701b514
Show file tree

Hide file tree

Showing 6 changed files with 48 additions and 8 deletions.
diff --git a/README.md b/README.md
@@ -76,6 +76,7 @@ List of all available options:
 | 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    |
+| selector             | The selector property lets you target a specific element to perform the screenshot on. (default `body`)    | string                    | optional    |
 
 ### Setting output image resolution
 

diff --git a/src/index.js b/src/index.js
@@ -8,6 +8,7 @@ module.exports = async function(options) {
     html,
     content,
     output,
+    selector = 'body',
     puppeteerArgs = {},
   } = options
 
@@ -23,23 +24,27 @@ module.exports = async function(options) {
 
   let buffers = []
 
-  await cluster.task(async ({ page, data: { content, output } }) => {
-    const buffer = await makeScreenshot(page, { ...options, content, output })
-
+  await cluster.task(async ({ page, data: { content, output, selector } }) => {
+    const buffer = await makeScreenshot(page, { ...options, content, output, selector })    
     buffers.push(buffer);
   });
 
+  cluster.on('taskerror', (err, data) => {
+    throw err
+  });
+
   const shouldBatch = Array.isArray(content)
-  const contents = shouldBatch ? content : [{ ...content, output }]
+  const contents = shouldBatch ? content : [{ ...content, output, selector }]
 
   contents.forEach(content => {
-    const { output, ...pageContent } = content
-    cluster.queue({ output, content: pageContent })
+    const { output, selector: contentSelector, ...pageContent } = content
+    cluster.queue({ output, content: pageContent, selector: contentSelector ? contentSelector : selector })
   })
 
   await cluster.idle();
   await cluster.close();
 
   return shouldBatch ? buffers : buffers[0]
+
 }
 
diff --git a/src/index.spec.js b/src/index.spec.js
@@ -70,6 +70,18 @@ describe('node-html-to-image', () => {
       const text = await getTextFromImage('./generated/image.png')
       expect(text.trim()).toBe('Hello Yvonnick!')
     })
+
+    it('should create selected element image', async () => {
+      await nodeHtmlToImage({
+        output: './generated/image.png',
+        html: '<html><body>Hello <div id="section">{{name}}!</div></body></html>',
+        content: { name: 'Sangwoo' },
+        selector: 'div#section',
+      })
+
+      const text = await getTextFromImage('./generated/image.png')
+      expect(text.trim()).toBe('Sangwoo!')
+    })
   })
 
   describe('batch', () => {
@@ -100,6 +112,18 @@ describe('node-html-to-image', () => {
       expect(result[1]).toBeInstanceOf(Buffer)
     })
 
+    it('should create selected elements images', async () => {
+      await nodeHtmlToImage({
+        html: '<html><body>Hello <div id="section1">{{name}}!</div><div id="section2">World!</div></body></html>',
+        content: [{ name: 'Sangwoo', output: './generated/image1.png', selector: 'div#section1' }, {output: './generated/image2.png', selector: 'div#section2'}],
+      })
+
+      const text1 = await getTextFromImage('./generated/image1.png')
+      expect(text1.trim()).toBe('Sangwoo!')
+      const text2 = await getTextFromImage('./generated/image2.png')
+      expect(text2.trim()).toBe('World!')
+    })
+
     it.skip('should handle mass volume well', async () => {
       jest.setTimeout(60000 * 60)
       expect.hasAssertions();

diff --git a/src/screenshot.js b/src/screenshot.js
@@ -8,6 +8,7 @@ module.exports = {
     encoding,
     content,
     html,
+    selector,
     beforeScreenshot,
     transparent = false,
     waitUntil = 'networkidle0',
@@ -22,7 +23,10 @@ module.exports = {
       html = template(content)
     }
     await page.setContent(html, { waitUntil })
-    const element = await page.$('body')
+    const element = await page.$(selector)
+    if (!element) {
+      throw Error('No element matches selector: ' + selector)
+    }
     if (beforeScreenshot && typeof beforeScreenshot === "function") {
       await beforeScreenshot(page);
     }

diff --git a/types/options.d.ts b/types/options.d.ts
@@ -50,4 +50,10 @@ export interface NodeHtmlToImageOptions {
    * @default binary
    */
   encoding?: 'binary' | 'base64';
+
+  /**
+   * The selector property lets you target a specific element to perform the screenshot on.
+   * @default body
+   */
+  selector?: string;
 }
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, beforeScreenshot, transparent, waitUntil }: NodeHtmlToImageOptions): Promise<string | Buffer>;
+declare function makeScreenshot(page: Page, { output, type, quality, encoding, content, html, beforeScreenshot, transparent, waitUntil, selector }: NodeHtmlToImageOptions): Promise<string | Buffer>;
 export default makeScreenshot;