cli: support png screenshots

Cargo.lock

@@ -284,6 +284,15 @@ version = "0.8.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "06ea2b9bc92be3c2baa9334a323ebca2d6f074ff852cd1d7b11064035cd3868f"
 
+[[package]]
+name = "crc32fast"
+version = "1.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3855a8a784b474f333699ef2bbca9db2c4a1f6d9088a90a2d25b1eb53111eaa"
+dependencies = [
+ "cfg-if",
+]
+
 [[package]]
 name = "deranged"
 version = "0.3.11"
@@ -356,6 +365,25 @@ version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "31ae425815400e5ed474178a7a22e275a9687086a12ca63ec793ff292d8fdae8"
 
+[[package]]
+name = "fdeflate"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4f9bfee30e4dedf0ab8b422f03af778d9612b63f502710fc500a334ebe2de645"
+dependencies = [
+ "simd-adler32",
+]
+
+[[package]]
+name = "flate2"
+version = "1.0.28"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "46303f565772937ffe1d394a4fac6f411c6013172fadde9dcdb1e147a086940e"
+dependencies = [
+ "crc32fast",
+ "miniz_oxide",
+]
+
 [[package]]
 name = "futures-core"
 version = "0.3.30"
@@ -553,6 +581,7 @@ dependencies = [
  "once_cell",
  "parking_lot",
  "pin-project",
+ "png",
  "rand",
  "repc",
  "serde",
@@ -682,6 +711,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9d811f3e15f28568be3407c8e7fdb6514c1cda3cb30683f15b6a1a1dc4ea14a7"
 dependencies = [
  "adler",
+ "simd-adler32",
 ]
 
 [[package]]
@@ -837,6 +867,19 @@ version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
 
+[[package]]
+name = "png"
+version = "0.17.13"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "06e4b0d3d1312775e782c86c91a111aa1f910cbb65e1337f9975b5f9a554b5e1"
+dependencies = [
+ "bitflags 1.3.2",
+ "crc32fast",
+ "fdeflate",
+ "flate2",
+ "miniz_oxide",
+]
+
 [[package]]
 name = "powerfmt"
 version = "0.2.0"
@@ -1076,6 +1119,12 @@ dependencies = [
  "roxmltree",
 ]
 
+[[package]]
+name = "simd-adler32"
+version = "0.3.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe"
+
 [[package]]
 name = "simplelog"
 version = "0.12.2"

Cargo.toml

@@ -56,6 +56,7 @@ gpu-alloc = "0.6.0"
 gpu-alloc-ash = "0.6.0"
 serde = { version = "1.0.196", features = ["derive"] }
 enum-map = "2.7.3"
+png = "0.17.13"
 
 [build-dependencies]
 repc = "0.1.1"

src/cli.rs

@@ -120,11 +120,23 @@ pub struct IdleSetArgs {
     pub interval: Vec<String>,
 }
 
+#[derive(ValueEnum, Debug, Copy, Clone, Hash, Default, PartialEq)]
+pub enum ScreenshotFormat {
+    /// The PNG image format.
+    #[default]
+    Png,
+    /// The QOI image format.
+    Qoi,
+}
+
 #[derive(Args, Debug)]
 pub struct ScreenshotArgs {
+    /// The format to use for the image.
+    #[clap(value_enum, long, default_value_t)]
+    pub format: ScreenshotFormat,
     /// The filename of the saved screenshot
     ///
-    /// If no filename is given, the screenshot will be saved under %Y-%m-%d-%H%M%S_jay.qoi
+    /// If no filename is given, the screenshot will be saved under %Y-%m-%d-%H%M%S_jay.<ext>
     /// in the current directory.
     ///
     /// The filename can contain the usual strftime parameters.

src/cli/screenshot.rs

@@ -1,6 +1,6 @@
 use {
     crate::{
-        cli::{GlobalArgs, ScreenshotArgs},
+        cli::{GlobalArgs, ScreenshotArgs, ScreenshotFormat},
         format::XRGB8888,
         tools::tool_client::{with_tool_client, Handle, ToolClient},
         utils::{errorfmt::ErrorFmt, queue::AsyncQueue},
@@ -16,6 +16,7 @@ use {
     },
     chrono::Local,
     jay_algorithms::qoi::xrgb8888_encode_qoi,
+    png::{BitDepth, ColorType, Encoder, SrgbRenderingIntent},
     std::rc::Rc,
 };
 
@@ -55,19 +56,25 @@ async fn run(screenshot: Rc<Screenshot>) {
             fatal!("Could not take a screenshot: {}", e);
         }
     };
-    let data = buf_to_qoi(&DmaBufIds::default(), &buf);
+    let format = screenshot.args.format;
-    let filename = screenshot
+    let data = buf_to_bytes(&DmaBufIds::default(), &buf, format);
-        .args
+    let filename = match &screenshot.args.filename {
-        .filename
+        Some(f) => f.clone(),
-        .as_deref()
+        _ => {
-        .unwrap_or("%Y-%m-%d-%H%M%S_jay.qoi");
+            let ext = match format {
-    let filename = Local::now().format(filename).to_string();
+                ScreenshotFormat::Png => "png",
+                ScreenshotFormat::Qoi => "qoi",
+            };
+            format!("%Y-%m-%d-%H%M%S_jay.{ext}")
+        }
+    };
+    let filename = Local::now().format(&filename).to_string();
     if let Err(e) = std::fs::write(&filename, data) {
         fatal!("Could not write `{}`: {}", filename, ErrorFmt(e));
     }
 }
 
-pub fn buf_to_qoi(dma_buf_ids: &DmaBufIds, buf: &Dmabuf) -> Vec<u8> {
+pub fn buf_to_bytes(dma_buf_ids: &DmaBufIds, buf: &Dmabuf, format: ScreenshotFormat) -> Vec<u8> {
     let drm = match Drm::reopen(buf.drm_dev.raw(), false) {
         Ok(drm) => drm,
         Err(e) => {
@@ -107,5 +114,22 @@ pub fn buf_to_qoi(dma_buf_ids: &DmaBufIds, buf: &Dmabuf) -> Vec<u8> {
         }
     };
     let data = unsafe { bo_map.data() };
-    xrgb8888_encode_qoi(data, buf.width, buf.height, buf.stride)
+    if format == ScreenshotFormat::Qoi {
+        return xrgb8888_encode_qoi(data, buf.width, buf.height, buf.stride);
+    }
+
+    let mut out = vec![];
+    {
+        let mut image_data = Vec::with_capacity(data.len());
+        for i in 0..data.len() / 4 {
+            image_data.extend_from_slice(&[data[4 * i + 2], data[4 * i + 1], data[4 * i + 0], 255])
+        }
+        let mut encoder = Encoder::new(&mut out, buf.width, buf.height);
+        encoder.set_color(ColorType::Rgba);
+        encoder.set_depth(BitDepth::Eight);
+        encoder.set_srgb(SrgbRenderingIntent::Perceptual);
+        let mut writer = encoder.write_header().unwrap();
+        writer.write_image_data(&image_data).unwrap();
+    }
+    out
 }

src/it/test_client.rs

@@ -1,6 +1,6 @@
 use {
     crate::{
-        cli::screenshot::buf_to_qoi,
+        cli::{screenshot::buf_to_bytes, ScreenshotFormat},
         client::Client,
         globals::GlobalBase,
         it::{
@@ -93,7 +93,11 @@ impl TestClient {
 
     pub async fn take_screenshot(&self, include_cursor: bool) -> Result<Vec<u8>, TestError> {
         let dmabuf = self.jc.take_screenshot(include_cursor).await?;
-        let qoi = buf_to_qoi(&self.server.state.dma_buf_ids, &dmabuf);
+        let qoi = buf_to_bytes(
+            &self.server.state.dma_buf_ids,
+            &dmabuf,
+            ScreenshotFormat::Qoi,
+        );
         Ok(qoi)
     }