Merge pull request #5 from YusukeIwaki/content_to_pdf

Implement Page#pdf
YusukeIwaki · Jun 28, 2020 · f510740 · f510740
2 parents e817e25 + 451a800
commit f510740
Show file tree

Hide file tree

Showing 4 changed files with 262 additions and 95 deletions.
diff --git a/lib/puppeteer/dom_world.rb b/lib/puppeteer/dom_world.rb
@@ -132,45 +132,48 @@ def SS(selector)
     document.SS(selector)
   end
 
-  # /**
-  #  * @return {!Promise<String>}
-  #  */
-  # async content() {
-  #   return await this.evaluate(() => {
-  #     let retVal = '';
-  #     if (document.doctype)
-  #       retVal = new XMLSerializer().serializeToString(document.doctype);
-  #     if (document.documentElement)
-  #       retVal += document.documentElement.outerHTML;
-  #     return retVal;
-  #   });
-  # }
+  # @return [String]
+  def content
+    evaluate <<-JAVASCRIPT
+    () => {
+      let retVal = '';
+      if (document.doctype)
+        retVal = new XMLSerializer().serializeToString(document.doctype);
+      if (document.documentElement)
+        retVal += document.documentElement.outerHTML;
+      return retVal;
+    }
+    JAVASCRIPT
+  end
 
-  # /**
-  #  * @param {string} html
-  #  * @param {!{timeout?: number, waitUntil?: string|!Array<string>}=} options
-  #  */
-  # async setContent(html, options = {}) {
-  #   const {
-  #     waitUntil = ['load'],
-  #     timeout = this._timeoutSettings.navigationTimeout(),
-  #   } = options;
-  #   // We rely upon the fact that document.open() will reset frame lifecycle with "init"
-  #   // lifecycle event. @see https://crrev.com/608658
-  #   await this.evaluate(html => {
-  #     document.open();
-  #     document.write(html);
-  #     document.close();
-  #   }, html);
-  #   const watcher = new LifecycleWatcher(this._frameManager, this._frame, waitUntil, timeout);
-  #   const error = await Promise.race([
-  #     watcher.timeoutOrTerminationPromise(),
-  #     watcher.lifecyclePromise(),
-  #   ]);
-  #   watcher.dispose();
-  #   if (error)
-  #     throw error;
-  # }
+  # @param html [String]
+  # @param timeout [Integer]
+  # @param wait_until [String|Array<String>]
+  def set_content(html, timeout: nil, wait_until: nil)
+    option_wait_until = [wait_until || 'load'].flatten
+    option_timeout = @timeout_settings.navigation_timeout
+
+    # We rely upon the fact that document.open() will reset frame lifecycle with "init"
+    # lifecycle event. @see https://crrev.com/608658
+    js = <<-JAVASCRIPT
+    (html) => {
+      document.open();
+      document.write(html);
+      document.close();
+    }
+    JAVASCRIPT
+    evaluate(js, html)
+
+    watcher = Puppeteer::LifecycleWatcher.new(@frame_manager, @frame, option_wait_until, option_timeout)
+    begin
+      await_any(
+        watcher.timeout_or_termination_promise,
+        watcher.lifecycle_promise,
+      )
+    ensure
+      watcher.dispose
+    end
+  end
 
   # /**
   #  * @param {!{url?: string, path?: string, content?: string, type?: string}} options

diff --git a/lib/puppeteer/frame.rb b/lib/puppeteer/frame.rb
@@ -110,12 +110,14 @@ def SS(selector)
 
   define_async_method :async_SS
 
+  # @return [String]
   def content
     @secondary_world.content
   end
 
-  # @param {string} html
-  # @param {!{timeout?: number, waitUntil?: string|!Array<string>}=} options
+  # @param html [String]
+  # @param timeout [Integer]
+  # @param wait_until [String|Array<String>]
   def set_content(html, timeout: nil, wait_until: nil)
     @secondary_world.set_content(html, timeout: timeout, wait_until: wait_until)
   end

diff --git a/lib/puppeteer/page.rb b/lib/puppeteer/page.rb
@@ -1,5 +1,7 @@
 require 'base64'
+require "stringio"
 
+require_relative './page/pdf_options'
 require_relative './page/screenshot_options'
 
 class Puppeteer::Page
@@ -649,13 +651,14 @@ def content
     main_frame.content
   end
 
-  # @param {string} html
-  # @param {!{timeout?: number, waitUntil?: string|!Array<string>}=} options
+  # @param html [String]
+  # @param timeout [Integer]
+  # @param wait_until [String|Array<String>]
   def set_content(html, timeout: nil, wait_until: nil)
     main_frame.set_content(html, timeout: timeout, wait_until: wait_until)
   end
 
-  # @param {string} html
+  # @param html [String]
   def content=(html)
     main_frame.set_content(html)
   end
@@ -915,60 +918,53 @@ def screenshot(options = {})
     buffer
   end
 
-  # /**
-  #  * @param {!PDFOptions=} options
-  #  * @return {!Promise<!Buffer>}
-  #  */
-  # async pdf(options = {}) {
-  #   const {
-  #     scale = 1,
-  #     displayHeaderFooter = false,
-  #     headerTemplate = '',
-  #     footerTemplate = '',
-  #     printBackground = false,
-  #     landscape = false,
-  #     pageRanges = '',
-  #     preferCSSPageSize = false,
-  #     margin = {},
-  #     path = null
-  #   } = options;
+  class ProtocolStreamReader
+    def initialize(client:, handle:, path:)
+      @client = client
+      @handle = handle
+      @path = path
+    end
 
-  #   let paperWidth = 8.5;
-  #   let paperHeight = 11;
-  #   if (options.format) {
-  #     const format = Page.PaperFormats[options.format.toLowerCase()];
-  #     assert(format, 'Unknown paper format: ' + options.format);
-  #     paperWidth = format.width;
-  #     paperHeight = format.height;
-  #   } else {
-  #     paperWidth = convertPrintParameterToInches(options.width) || paperWidth;
-  #     paperHeight = convertPrintParameterToInches(options.height) || paperHeight;
-  #   }
+    def read
+      out = StringIO.new
+      File.open(@path, 'w') do |file|
+        eof = false
+        until eof
+          response = @client.send_message('IO.read', handle: @handle)
+          eof = response['eof']
+          data =
+            if response['base64Encoded']
+              Base64.decode64(response['data'])
+            else
+              response['data']
+            end
+          out.write(data)
+          file.write(data)
+        end
+      end
+      @client.send_message('IO.close', handle: @handle)
+      out.read
+    end
+  end
 
-  #   const marginTop = convertPrintParameterToInches(margin.top) || 0;
-  #   const marginLeft = convertPrintParameterToInches(margin.left) || 0;
-  #   const marginBottom = convertPrintParameterToInches(margin.bottom) || 0;
-  #   const marginRight = convertPrintParameterToInches(margin.right) || 0;
-
-  #   const result = await this._client.send('Page.printToPDF', {
-  #     transferMode: 'ReturnAsStream',
-  #     landscape,
-  #     displayHeaderFooter,
-  #     headerTemplate,
-  #     footerTemplate,
-  #     printBackground,
-  #     scale,
-  #     paperWidth,
-  #     paperHeight,
-  #     marginTop,
-  #     marginBottom,
-  #     marginLeft,
-  #     marginRight,
-  #     pageRanges,
-  #     preferCSSPageSize
-  #   });
-  #   return await helper.readProtocolStream(this._client, result.stream, path);
-  # }
+  class PrintToPdfIsNotImplementedError < StandardError
+    def initialize
+      super('pdf() is only available in headless mode. See https://github.com/puppeteer/puppeteer/issues/1829')
+    end
+  end
+
+  # @return [String]
+  def pdf(options = {})
+    pdf_options = PDFOptions.new(options)
+    result = @client.send_message('Page.printToPDF', pdf_options.page_print_args)
+    ProtocolStreamReader.new(client: @client, handle: result['stream'], path: pdf_options.path).read
+  rescue => err
+    if err.message.include?('PrintToPDF is not implemented')
+      raise PrintToPdfIsNotImplementedError.new
+    else
+      raise
+    end
+  end
 
   # @param run_before_unload [Boolean]
   def close(run_before_unload: false)