webshot 0.3.0

A command-line tool for automated website screenshots and web scraping
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
use crate::error::{Result, WebshotError};
use std::path::Path;

/// Scrolling screenshot mode
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum ScrollMode {
    /// Standard viewport-only screenshot
    Viewport,
    /// Capture full page by scrolling and stitching
    FullPage,
    /// Capture element and its full scrollable content
    FullElement,
}

impl Default for ScrollMode {
    fn default() -> Self {
        Self::Viewport
    }
}

/// Screenshot configuration options
#[derive(Debug, Clone)]
pub struct ScreenshotOptions {
    /// Viewport width
    pub width: u32,
    /// Viewport height
    pub height: u32,
    /// CSS selector for element screenshot
    pub selector: Option<String>,
    /// JavaScript to execute before screenshot
    pub javascript: Option<String>,
    /// Element to wait for before taking screenshot
    pub wait_for: Option<String>,
    /// Timeout in seconds
    pub timeout: u64,
    /// Enable retina/high-DPI mode
    pub retina: bool,
    /// JPEG quality (1-100)
    pub quality: Option<u8>,
    /// Wait time before taking screenshot
    pub wait: u64,
    /// Custom user agent
    pub user_agent: Option<String>,
    /// Scrolling screenshot mode
    pub scroll_mode: ScrollMode,
    /// Maximum height for full page screenshots (safety limit)
    pub max_height: Option<u32>,
    /// Scroll delay between captures (milliseconds)
    pub scroll_delay: u64,
}

impl Default for ScreenshotOptions {
    fn default() -> Self {
        Self {
            width: 1280,
            height: 800,
            selector: None,
            javascript: None,
            wait_for: None,
            timeout: 30,
            retina: false,
            quality: None,
            wait: 0,
            user_agent: None,
            scroll_mode: ScrollMode::default(),
            max_height: Some(30000), // 30k pixels default max height
            scroll_delay: 100,       // 100ms delay between scrolls
        }
    }
}

impl ScreenshotOptions {
    /// Create new screenshot options with default values
    pub fn new() -> Self {
        Self::default()
    }

    /// Set viewport dimensions
    pub fn viewport(mut self, width: u32, height: u32) -> Self {
        self.width = width;
        self.height = height;
        self
    }

    /// Set CSS selector for element screenshot
    pub fn selector<S: Into<String>>(mut self, selector: S) -> Self {
        self.selector = Some(selector.into());
        self
    }

    /// Set JavaScript to execute
    pub fn javascript<S: Into<String>>(mut self, script: S) -> Self {
        self.javascript = Some(script.into());
        self
    }

    /// Set element to wait for
    pub fn wait_for<S: Into<String>>(mut self, selector: S) -> Self {
        self.wait_for = Some(selector.into());
        self
    }

    /// Set timeout in seconds
    pub fn timeout(mut self, timeout: u64) -> Self {
        self.timeout = timeout;
        self
    }

    /// Enable retina mode
    pub fn retina(mut self) -> Self {
        self.retina = true;
        self
    }

    /// Set JPEG quality
    pub fn quality(mut self, quality: u8) -> Self {
        self.quality = Some(quality);
        self
    }

    /// Set wait time before screenshot
    pub fn wait(mut self, wait: u64) -> Self {
        self.wait = wait;
        self
    }

    /// Set custom user agent
    pub fn user_agent<S: Into<String>>(mut self, user_agent: S) -> Self {
        self.user_agent = Some(user_agent.into());
        self
    }

    /// Enable full page scrolling screenshot
    pub fn full_page(mut self) -> Self {
        self.scroll_mode = ScrollMode::FullPage;
        self
    }

    /// Enable full element scrolling screenshot
    pub fn full_element(mut self) -> Self {
        self.scroll_mode = ScrollMode::FullElement;
        self
    }

    /// Set scroll mode
    pub fn scroll_mode(mut self, mode: ScrollMode) -> Self {
        self.scroll_mode = mode;
        self
    }

    /// Set maximum height for full page screenshots
    pub fn max_height(mut self, height: u32) -> Self {
        self.max_height = Some(height);
        self
    }

    /// Set scroll delay between captures
    pub fn scroll_delay(mut self, delay: u64) -> Self {
        self.scroll_delay = delay;
        self
    }

    /// Validate the options
    pub fn validate(&self) -> Result<()> {
        if self.width == 0 || self.height == 0 {
            return Err(WebshotError::InvalidViewport {
                width: self.width,
                height: self.height,
            });
        }

        if let Some(quality) = self.quality {
            if !(1..=100).contains(&quality) {
                return Err(WebshotError::config(format!(
                    "JPEG quality must be between 1-100, got: {}",
                    quality
                )));
            }
        }

        if self.timeout == 0 {
            return Err(WebshotError::config(
                "Timeout must be greater than 0".to_string(),
            ));
        }

        // Validate max height for full page screenshots
        if let Some(max_height) = self.max_height {
            if max_height == 0 {
                return Err(WebshotError::config(
                    "Max height must be greater than 0".to_string(),
                ));
            }
            if max_height > 100000 {
                return Err(WebshotError::config(
                    "Max height should not exceed 100,000 pixels for safety".to_string(),
                ));
            }
        }

        // Validate scroll mode compatibility
        if self.scroll_mode == ScrollMode::FullElement && self.selector.is_none() {
            return Err(WebshotError::config(
                "FullElement scroll mode requires a selector to be specified".to_string(),
            ));
        }

        Ok(())
    }

    /// Get device scale factor based on retina setting
    pub fn device_scale_factor(&self) -> f64 {
        if self.retina {
            2.0
        } else {
            1.0
        }
    }

    /// Determine output format from file path
    pub fn output_format<P: AsRef<Path>>(&self, path: P) -> Result<ImageFormat> {
        let extension = path
            .as_ref()
            .extension()
            .and_then(|ext| ext.to_str())
            .map(|ext| ext.to_lowercase())
            .ok_or_else(|| WebshotError::config("No file extension found".to_string()))?;

        match extension.as_str() {
            "png" => Ok(ImageFormat::Png),
            "jpg" | "jpeg" => Ok(ImageFormat::Jpeg),
            "pdf" => Ok(ImageFormat::Pdf),
            "webp" => Ok(ImageFormat::WebP),
            _ => Err(WebshotError::UnsupportedFormat { format: extension }),
        }
    }
}

/// Supported image formats
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ImageFormat {
    Png,
    Jpeg,
    WebP,
    Pdf,
}

impl ImageFormat {
    /// Get the default file extension for this format
    pub fn extension(&self) -> &'static str {
        match self {
            ImageFormat::Png => "png",
            ImageFormat::Jpeg => "jpg",
            ImageFormat::WebP => "webp",
            ImageFormat::Pdf => "pdf",
        }
    }

    /// Get the MIME type for this format
    pub fn mime_type(&self) -> &'static str {
        match self {
            ImageFormat::Png => "image/png",
            ImageFormat::Jpeg => "image/jpeg",
            ImageFormat::WebP => "image/webp",
            ImageFormat::Pdf => "application/pdf",
        }
    }

    /// Check if this format supports quality settings
    pub fn supports_quality(&self) -> bool {
        matches!(self, ImageFormat::Jpeg | ImageFormat::WebP)
    }

    /// Check if this format supports transparency
    pub fn supports_transparency(&self) -> bool {
        matches!(self, ImageFormat::Png | ImageFormat::WebP)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;

    #[test]
    fn test_screenshot_options_builder() {
        let options = ScreenshotOptions::new()
            .viewport(1920, 1080)
            .selector(".header")
            .javascript("console.log('test')")
            .wait_for(".content")
            .timeout(60)
            .retina()
            .quality(90)
            .wait(5)
            .user_agent("Custom Agent")
            .full_page()
            .max_height(50000)
            .scroll_delay(200);

        assert_eq!(options.width, 1920);
        assert_eq!(options.height, 1080);
        assert_eq!(options.selector.as_deref(), Some(".header"));
        assert_eq!(options.javascript.as_deref(), Some("console.log('test')"));
        assert_eq!(options.wait_for.as_deref(), Some(".content"));
        assert_eq!(options.timeout, 60);
        assert!(options.retina);
        assert_eq!(options.quality, Some(90));
        assert_eq!(options.wait, 5);
        assert_eq!(options.user_agent.as_deref(), Some("Custom Agent"));
        assert_eq!(options.scroll_mode, ScrollMode::FullPage);
        assert_eq!(options.max_height, Some(50000));
        assert_eq!(options.scroll_delay, 200);
    }

    #[test]
    fn test_options_validation() {
        let mut options = ScreenshotOptions::new();
        assert!(options.validate().is_ok());

        options.width = 0;
        assert!(options.validate().is_err());

        options.width = 1280;
        options.quality = Some(150);
        assert!(options.validate().is_err());

        options.quality = Some(80);
        options.timeout = 0;
        assert!(options.validate().is_err());
    }

    #[test]
    fn test_output_format_detection() {
        let options = ScreenshotOptions::new();

        assert_eq!(
            options.output_format(PathBuf::from("test.png")).unwrap(),
            ImageFormat::Png
        );
        assert_eq!(
            options.output_format(PathBuf::from("test.jpg")).unwrap(),
            ImageFormat::Jpeg
        );
        assert_eq!(
            options.output_format(PathBuf::from("test.jpeg")).unwrap(),
            ImageFormat::Jpeg
        );
        assert_eq!(
            options.output_format(PathBuf::from("test.pdf")).unwrap(),
            ImageFormat::Pdf
        );
        assert_eq!(
            options.output_format(PathBuf::from("test.webp")).unwrap(),
            ImageFormat::WebP
        );

        assert!(options.output_format(PathBuf::from("test.gif")).is_err());
        assert!(options.output_format(PathBuf::from("test")).is_err());
    }

    #[test]
    fn test_device_scale_factor() {
        let options = ScreenshotOptions::new();
        assert_eq!(options.device_scale_factor(), 1.0);

        let retina_options = options.retina();
        assert_eq!(retina_options.device_scale_factor(), 2.0);
    }

    #[test]
    fn test_image_format() {
        assert_eq!(ImageFormat::Png.extension(), "png");
        assert_eq!(ImageFormat::Jpeg.extension(), "jpg");
        assert_eq!(ImageFormat::Pdf.extension(), "pdf");
        assert_eq!(ImageFormat::WebP.extension(), "webp");

        assert_eq!(ImageFormat::Png.mime_type(), "image/png");
        assert_eq!(ImageFormat::Jpeg.mime_type(), "image/jpeg");
        assert_eq!(ImageFormat::Pdf.mime_type(), "application/pdf");
        assert_eq!(ImageFormat::WebP.mime_type(), "image/webp");

        assert!(!ImageFormat::Png.supports_quality());
        assert!(ImageFormat::Jpeg.supports_quality());
        assert!(!ImageFormat::Pdf.supports_quality());
        assert!(ImageFormat::WebP.supports_quality());

        assert!(ImageFormat::Png.supports_transparency());
        assert!(!ImageFormat::Jpeg.supports_transparency());
        assert!(!ImageFormat::Pdf.supports_transparency());
        assert!(ImageFormat::WebP.supports_transparency());
    }

    #[test]
    fn test_scroll_mode_defaults() {
        let options = ScreenshotOptions::new();
        assert_eq!(options.scroll_mode, ScrollMode::Viewport);
        assert_eq!(options.max_height, Some(30000));
        assert_eq!(options.scroll_delay, 100);
    }

    #[test]
    fn test_scroll_mode_builder() {
        let full_page_options = ScreenshotOptions::new().full_page();
        assert_eq!(full_page_options.scroll_mode, ScrollMode::FullPage);

        let full_element_options = ScreenshotOptions::new().selector(".content").full_element();
        assert_eq!(full_element_options.scroll_mode, ScrollMode::FullElement);
    }

    #[test]
    fn test_scroll_validation() {
        // Full element mode requires a selector
        let mut options = ScreenshotOptions::new().full_element();
        assert!(options.validate().is_err());

        // With selector it should work
        options = options.selector(".content");
        assert!(options.validate().is_ok());

        // Max height validation
        options = ScreenshotOptions::new().max_height(0);
        assert!(options.validate().is_err());

        options = ScreenshotOptions::new().max_height(200000); // Too large
        assert!(options.validate().is_err());
    }
}