Merge pull request #176 from mahkoh/jorth/png

cli: support png screenshots

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 filename = screenshot
-        .args
-        .filename
-        .as_deref()
-        .unwrap_or("%Y-%m-%d-%H%M%S_jay.qoi");
-    let filename = Local::now().format(filename).to_string();
+    let format = screenshot.args.format;
+    let data = buf_to_bytes(&DmaBufIds::default(), &buf, format);
+    let filename = match &screenshot.args.filename {
+        Some(f) => f.clone(),
+        _ => {
+            let ext = match format {
+                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)
     }