webscreenshotlib/
lib.rs

1use headless_chrome::{Browser, protocol::{target::methods::CreateTarget ,page::ScreenshotFormat}};
2
3/// Formats available to encode screenshots to.
4pub enum OutputFormat {PNG, JPG, PDF}
5
6/// Take a screenshot of a webpage rendered in a Chrome tab.
7/// 
8/// ### Arguments
9/// - url: Web or file URL to render for screenshot.
10/// - format: The output format to encode the screenshot as.
11/// - quality: JPG quality. 0-100. This is ignored for PNGs and PDFs.
12/// - surface: If true, screencap the entire rended HTML. If false, screencap what is visible inside the browser viewport dimensions.
13/// - width: Width of the Chrome browser.
14/// - height: Height of the Chrome browser.
15/// - element: CSS selector to screenshot instead of the full page. Empty string screenshots whole page.
16///
17/// ### Example
18/// ```
19/// # use webscreenshotlib::{OutputFormat, screenshot_tab};
20/// # std::fs::write("/tmp/test.html", "<html><head><title>Title</title></head><body>Hello World</body></html>");
21/// # fn test_screenshot_tab() -> Result<(), failure::Error> {
22/// // Screenshot a page.
23/// let image_data = screenshot_tab(
24///   "file:////tmp/test.html",
25///   OutputFormat::PNG, 80, false, 1024, 800, "")?;
26/// # more_asserts::assert_ge!(image_data.len(), 1024*800);
27/// # return Ok(())
28/// # }
29/// ```
30// TODO: Add element screenshot feature.
31pub fn screenshot_tab(
32    url: &str,
33    format: OutputFormat,
34    quality: u8,
35    visible_only: bool,
36    width: u16,
37    height: u16,
38    element: &str)
39    -> Result<Vec<u8>, failure::Error>
40{
41    // Get browser, navigate to page.
42    let browser = Browser::default()?;
43    let tab = browser.new_tab_with_options(CreateTarget {
44        url: url,
45        width: Some(width.into()),
46        height: Some(height.into()),
47        browser_context_id: None,
48        enable_begin_frame_control: None,
49    })?;
50    tab.navigate_to(url)?;
51    tab.wait_until_navigated()?;
52
53    // If element arg is supplied, wait for it to load. This will signal the
54    // tab.capture_screenshot method to screenshot only the element.
55    if element.is_empty() != true {
56        tab.wait_for_element(element)?;
57    }
58
59    // Take screenshot in given format, return image bytes.
60    match format {
61        OutputFormat::JPG => {
62            return Ok(tab.capture_screenshot(
63                ScreenshotFormat::JPEG(Some(quality.into())), None, !visible_only)?);
64            },
65        OutputFormat::PNG => return Ok(tab.capture_screenshot(
66            ScreenshotFormat::PNG, None, !visible_only)?),
67        OutputFormat::PDF => return Ok(
68            tab.print_to_pdf(None)?),
69    };
70}
71
72/// Write image data to a file. A simple wrapper around std::fs::write().
73pub fn write_screenshot(
74    output_path: &str,
75    img_data: Vec<u8>)
76    -> Result<(), failure::Error>
77{
78    std::fs::write(output_path, img_data)?;
79    return Ok(())
80}