From 6410b68cd201a17f593a7a8f7cc2605b138d963a Mon Sep 17 00:00:00 2001
From: Erick Guan <erickguan@users.noreply.github.com>
Date: Thu, 4 Jul 2024 19:07:47 +0200
Subject: [PATCH] Support webp in screenshot

---
 CHANGELOG.md                  |  1 +
 README.md                     |  2 +-
 lib/ferrum/page/screenshot.rb | 31 +++++++++++++++++++++++--------
 spec/page/screenshot_spec.rb  | 11 +++++++++++
 4 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8d2c9a5d..93383dc1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@
 - `Ferrum::Node#exists?` check whether the node in ruby world still exists in the DOM tree.
 - `Ferrum::Cookies#store` stores all cookies of current page in a file.
 - `Ferrum::Cookies#load` Loads all cookies from the file and sets them for current page.
+- `Ferrum::Page#screenshot` supports webp image format
 
 ### Changed
 
diff --git a/README.md b/README.md
index 893be91d..5d9e051e 100644
--- a/README.md
+++ b/README.md
@@ -398,7 +398,7 @@ Saves screenshot on a disk or returns it as base64.
     `:binary` automatically
   * :encoding `Symbol` `:base64` | `:binary` you can set it to return image as
     Base64
-  * :format `String` "jpeg" | "png"
+  * :format `String` "jpeg" ("jpg") | "png" | "webp"
   * :quality `Integer` 0-100 works for jpeg only
   * :full `Boolean` whether you need full page screenshot or a viewport
   * :selector `String` css selector for given element, optional
diff --git a/lib/ferrum/page/screenshot.rb b/lib/ferrum/page/screenshot.rb
index 94114d09..f30e7975 100644
--- a/lib/ferrum/page/screenshot.rb
+++ b/lib/ferrum/page/screenshot.rb
@@ -8,6 +8,9 @@ module Screenshot
       FULL_WARNING = "Ignoring :selector or :area in #screenshot since full: true was given at %s"
       AREA_WARNING = "Ignoring :area in #screenshot since selector: was given at %s"
 
+      DEFAULT_SCREENSHOT_FORMAT = "png"
+      SUPPORTED_SCREENSHOT_FORMAT = %w[png jpeg jpg webp].freeze
+
       DEFAULT_PDF_OPTIONS = {
         landscape: false,
         paper_width: 8.5,
@@ -41,7 +44,7 @@ module Screenshot
       # @option opts [:base64, :binary] :encoding
       #   The encoding the image should be returned in.
       #
-      # @option opts ["jpeg", "png"] :format
+      # @option opts ["jpeg", "jpg", "png", "webp"] :format
       #   The format the image should be returned in.
       #
       # @option opts [Integer] :quality
@@ -71,8 +74,11 @@ module Screenshot
       # @example Save on the disk in JPG:
       #   page.screenshot(path: "google.jpg") # => 30902
       #
+      # @example Save to Base64 in WebP with reduce quality:
+      #   page.screenshot(format: 'webp', quality: 60) # "iVBORw0KGgoAAAANS...
+      #
       # @example Save to Base64 the whole page not only viewport and reduce quality:
-      #   page.screenshot(full: true, quality: 60) # "iVBORw0KGgoAAAANS...
+      #   page.screenshot(full: true, format: 'jpeg', quality: 60) # "iVBORw0KGgoAAAANS...
       #
       # @example Save with specific background color:
       #   page.screenshot(background_color: Ferrum::RGBA.new(0, 0, 0, 0.0))
@@ -210,14 +216,23 @@ def screenshot_options(path = nil, format: nil, scale: 1.0, **options)
         screenshot_options
       end
 
-      def format_options(format, path, quality)
-        format ||= path ? File.extname(path).delete(".") : "png"
-        format = "jpeg" if format == "jpg"
-        raise "Not supported options `:format` #{format}. jpeg | png" if format !~ /jpeg|png/i
+      def format_options(screenshot_format, path, quality)
+        if !screenshot_format && path # try to infer from path
+          extension = File.extname(path).delete(".")&.downcase
+          screenshot_format = extension if extension && !extension.empty?
+        end
+
+        screenshot_format ||= DEFAULT_SCREENSHOT_FORMAT
+        unless SUPPORTED_SCREENSHOT_FORMAT.include?(screenshot_format)
+          raise "Not supported options `:format` #{screenshot_format}. #{SUPPORTED_SCREENSHOT_FORMAT.join(' | ')}"
+        end
+
+        screenshot_format = "jpeg" if screenshot_format == "jpg"
 
-        quality ||= 75 if format == "jpeg"
+        # Chrome supports screenshot qualities for JPEG and WebP
+        quality ||= 75 if screenshot_format != "png"
 
-        [format, quality]
+        [screenshot_format, quality]
       end
 
       def area_options(full, selector, scale, area = nil)
diff --git a/spec/page/screenshot_spec.rb b/spec/page/screenshot_spec.rb
index aebd8674..a55cbb8c 100644
--- a/spec/page/screenshot_spec.rb
+++ b/spec/page/screenshot_spec.rb
@@ -152,6 +152,17 @@ def create_screenshot(**options)
       FileUtils.rm_f(file)
     end
 
+    it "supports screenshotting the page to webp file" do
+      file = "#{PROJECT_ROOT}/spec/tmp/screenshot.webp"
+      browser.go_to
+
+      browser.screenshot(path: file, format: "webp")
+
+      expect(File.exist?(file)).to be true
+    ensure
+      FileUtils.rm_f(file)
+    end
+
     it "supports screenshotting the page with different quality settings" do
       file2 = "#{PROJECT_ROOT}/spec/tmp/screenshot2.jpeg"
       file3 = "#{PROJECT_ROOT}/spec/tmp/screenshot3.jpeg"