firefox_webdriver/browser/tab/
screenshot.rs

1//! Screenshot capture methods.
2
3use std::path::Path;
4
5use base64::Engine;
6use base64::engine::general_purpose::STANDARD as Base64Standard;
7use tracing::debug;
8
9use crate::error::{Error, Result};
10use crate::protocol::command::{BrowsingContextCommand, Command};
11
12use super::Tab;
13
14// ============================================================================
15// Types
16// ============================================================================
17
18/// Image format for screenshots.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
20pub enum ImageFormat {
21    /// PNG format (lossless, larger file size).
22    #[default]
23    Png,
24    /// JPEG format with quality (0-100).
25    Jpeg(u8),
26}
27
28impl ImageFormat {
29    /// Creates PNG format.
30    #[inline]
31    #[must_use]
32    pub fn png() -> Self {
33        Self::Png
34    }
35
36    /// Creates JPEG format with quality (0-100).
37    #[inline]
38    #[must_use]
39    pub fn jpeg(quality: u8) -> Self {
40        Self::Jpeg(quality.min(100))
41    }
42
43    /// Returns the MIME type for this format.
44    #[must_use]
45    pub fn mime_type(&self) -> &'static str {
46        match self {
47            Self::Png => "image/png",
48            Self::Jpeg(_) => "image/jpeg",
49        }
50    }
51
52    /// Returns the file extension for this format.
53    #[must_use]
54    pub fn extension(&self) -> &'static str {
55        match self {
56            Self::Png => "png",
57            Self::Jpeg(_) => "jpg",
58        }
59    }
60
61    /// Returns the format string for the protocol.
62    fn format_str(&self) -> &'static str {
63        match self {
64            Self::Png => "png",
65            Self::Jpeg(_) => "jpeg",
66        }
67    }
68
69    /// Returns the quality value if JPEG.
70    fn quality(&self) -> Option<u8> {
71        match self {
72            Self::Png => None,
73            Self::Jpeg(q) => Some(*q),
74        }
75    }
76}
77
78// ============================================================================
79// ScreenshotBuilder
80// ============================================================================
81
82/// Builder for configuring and capturing screenshots.
83///
84/// Uses the browser's native screenshot API (`browser.tabs.captureVisibleTab`)
85/// for accurate pixel capture without JavaScript limitations.
86///
87/// # Example
88///
89/// ```ignore
90/// // Capture as PNG base64
91/// let png_data = tab.screenshot().png().capture().await?;
92///
93/// // Capture as JPEG and save to file
94/// tab.screenshot().jpeg(80).save("page.jpg").await?;
95/// ```
96pub struct ScreenshotBuilder<'a> {
97    tab: &'a Tab,
98    format: ImageFormat,
99}
100
101impl<'a> ScreenshotBuilder<'a> {
102    /// Creates a new screenshot builder.
103    pub(crate) fn new(tab: &'a Tab) -> Self {
104        Self {
105            tab,
106            format: ImageFormat::Png,
107        }
108    }
109
110    /// Sets PNG format (default).
111    #[must_use]
112    pub fn png(mut self) -> Self {
113        self.format = ImageFormat::Png;
114        self
115    }
116
117    /// Sets JPEG format with quality (0-100).
118    #[must_use]
119    pub fn jpeg(mut self, quality: u8) -> Self {
120        self.format = ImageFormat::Jpeg(quality.min(100));
121        self
122    }
123
124    /// Sets the image format.
125    #[must_use]
126    pub fn format(mut self, format: ImageFormat) -> Self {
127        self.format = format;
128        self
129    }
130
131    /// Captures the screenshot and returns base64-encoded data.
132    ///
133    /// Uses the browser's native `captureVisibleTab` API for accurate capture.
134    pub async fn capture(&self) -> Result<String> {
135        debug!(
136            tab_id = %self.tab.inner.tab_id,
137            format = ?self.format,
138            "Capturing screenshot via browser API"
139        );
140
141        let command = Command::BrowsingContext(BrowsingContextCommand::CaptureScreenshot {
142            format: self.format.format_str().to_string(),
143            quality: self.format.quality(),
144        });
145
146        let response = self.tab.send_command(command).await?;
147
148        debug!(response = ?response, "Screenshot response");
149
150        let data = response
151            .result
152            .as_ref()
153            .and_then(|v| v.get("data"))
154            .and_then(|v| v.as_str())
155            .ok_or_else(|| {
156                let result_str = response
157                    .result
158                    .as_ref()
159                    .map(|v| v.to_string())
160                    .unwrap_or_else(|| "null".to_string());
161                Error::script_error(format!(
162                    "Screenshot response missing data field. Got: {}",
163                    result_str
164                ))
165            })?;
166
167        Ok(data.to_string())
168    }
169
170    /// Captures the screenshot and returns raw bytes.
171    pub async fn capture_bytes(&self) -> Result<Vec<u8>> {
172        let base64_data = self.capture().await?;
173        Base64Standard
174            .decode(&base64_data)
175            .map_err(|e| Error::script_error(format!("Failed to decode base64: {}", e)))
176    }
177
178    /// Captures the screenshot and saves to a file.
179    ///
180    /// The format is determined by the builder settings, not the file extension.
181    pub async fn save(&self, path: impl AsRef<Path>) -> Result<()> {
182        let bytes = self.capture_bytes().await?;
183        std::fs::write(path.as_ref(), bytes).map_err(Error::Io)?;
184        Ok(())
185    }
186}
187
188// ============================================================================
189// Tab - Screenshot
190// ============================================================================
191
192impl Tab {
193    /// Creates a screenshot builder for capturing page screenshots.
194    ///
195    /// Uses the browser's native screenshot API for accurate pixel capture.
196    ///
197    /// # Example
198    ///
199    /// ```ignore
200    /// // PNG screenshot as base64
201    /// let data = tab.screenshot().png().capture().await?;
202    ///
203    /// // JPEG screenshot saved to file
204    /// tab.screenshot().jpeg(85).save("page.jpg").await?;
205    /// ```
206    #[must_use]
207    pub fn screenshot(&self) -> ScreenshotBuilder<'_> {
208        ScreenshotBuilder::new(self)
209    }
210
211    /// Captures a PNG screenshot and returns base64-encoded data.
212    ///
213    /// Shorthand for `tab.screenshot().png().capture().await`.
214    pub async fn capture_screenshot(&self) -> Result<String> {
215        self.screenshot().png().capture().await
216    }
217
218    /// Captures a screenshot and saves to a file.
219    ///
220    /// Format is determined by file extension (.png or .jpg/.jpeg).
221    pub async fn save_screenshot(&self, path: impl AsRef<Path>) -> Result<()> {
222        let path = path.as_ref();
223        let ext = path
224            .extension()
225            .and_then(|e| e.to_str())
226            .unwrap_or("png")
227            .to_lowercase();
228
229        let builder = self.screenshot();
230        let builder = match ext.as_str() {
231            "jpg" | "jpeg" => builder.jpeg(85),
232            _ => builder.png(),
233        };
234
235        builder.save(path).await
236    }
237}
firefox_webdriver/browser/tab/screenshot.rs

firefox_webdriver/browser/tab/
screenshot.rs
