[backend.rs] Add support for ppm rendering.

* Closes: #12 * Thanks to AmusedAstronaut in our matrix room for patching ppm support in * Fixes a regression where `--stdout` would not respect `--extension` provided by user Signed-off-by: Shinyzenith <[email protected]>
waycrate · Jul 20, 2022 · bbbdc2c · bbbdc2c
1 parent 9a1f277
commit bbbdc2c
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 52 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -8,7 +8,7 @@ keywords = ["screenshot", "wayland", "wlroots"]
 license = "BSD-2-Clause"
 name = "wayshot"
 repository = "https://git.sr.ht/~shinyzenith/wayshot"
-version = "1.1.8"
+version = "1.1.9"
 exclude = [
     "CODE_OF_CONDUCT.md",
     "CONTRIBUTING.md",
@@ -20,7 +20,7 @@ exclude = [
 [dependencies]
 clap = "3.1.18"
 env_logger = { version = "0.9.0", default-features = false, features = ["atty", "termcolor"] }
-image = { version = "0.24", default-features = false, features = ["jpeg", "png"] }
+image = { version = "0.24", default-features = false, features = ["jpeg", "png", "pnm"] }
 log = "0.4.17"
 memmap2 = "0.5.3"
 nix = "0.24.1"

diff --git a/src/backend.rs b/src/backend.rs
@@ -19,7 +19,11 @@ use nix::{
 };
 
 use image::{
-    codecs::{jpeg::JpegEncoder, png::PngEncoder},
+    codecs::{
+        jpeg::JpegEncoder,
+        png::PngEncoder,
+        pnm::{self, PnmEncoder},
+    },
     ColorType,
     ColorType::Rgba8,
     ImageEncoder,
@@ -84,6 +88,8 @@ pub enum EncodingFormat {
     Jpg,
     /// Png encoder.
     Png,
+    /// Ppm encoder
+    Ppm,
 }
 
 /// Get a FrameCopy instance with screenshot pixel data for any wl_output object.
@@ -375,6 +381,28 @@ pub fn write_to_file(
             )?;
             let _ = output_file.flush()?;
         }
+        EncodingFormat::Ppm => {
+            let rgb8_data = if let ColorType::Rgba8 = frame_copy.frame_color_type {
+                let data = frame_copy
+                    .frame_mmap
+                    .iter()
+                    .enumerate()
+                    .filter_map(|(i, &v)| if i % 4 != 3 { Some(v) } else { None })
+                    .collect::<Vec<u8>>();
+                data
+            } else {
+                unimplemented!("Currently only ColorType::Rgba8 is supported")
+            };
+            PnmEncoder::new(&mut output_file)
+                .with_subtype(pnm::PnmSubtype::Pixmap(pnm::SampleEncoding::Binary))
+                .write_image(
+                    &rgb8_data,
+                    frame_copy.frame_format.width,
+                    frame_copy.frame_format.height,
+                    ColorType::Rgb8,
+                )?;
+            let _ = output_file.flush()?;
+        }
     }
 
     Ok(())

diff --git a/src/wayshot.rs b/src/wayshot.rs
@@ -29,7 +29,6 @@ fn main() -> Result<(), Box<dyn Error>> {
     log::trace!("Logger initialized.");
 
     let display = Display::connect_to_env()?;
-    let mut extension = backend::EncodingFormat::Png;
 
     let mut cursor_overlay: i32 = 0;
     if args.is_present("cursor") {
@@ -111,14 +110,32 @@ fn main() -> Result<(), Box<dyn Error>> {
         backend::capture_output_frame(display, cursor_overlay, output, None)?
     };
 
+    let extension = if args.is_present("extension") {
+        let ext = args.value_of("extension").unwrap().trim();
+        match ext {
+            "jpeg" | "jpg" => backend::EncodingFormat::Jpg,
+            "png" => backend::EncodingFormat::Png,
+            "ppm" => backend::EncodingFormat::Ppm,
+            _ => {
+                log::error!("Invalid extension provided.\nValid extensions:\n1) jpeg\n2) jpg\n3) png\n4) ppm");
+                exit(1);
+            }
+        }
+    } else {
+        backend::EncodingFormat::Png
+    };
+
+    if extension != backend::EncodingFormat::Png {
+        log::debug!("Using custom extension: {:#?}", extension);
+    }
+
     if args.is_present("stdout") {
         let stdout = stdout();
         let writer = BufWriter::new(stdout.lock());
         backend::write_to_file(writer, extension, frame_copy)?;
     } else {
-        let path: String;
-        if args.is_present("file") {
-            path = args.value_of("file").unwrap().trim().to_string();
+        let path = if args.is_present("file") {
+            args.value_of("file").unwrap().trim().to_string()
         } else {
             let time = match SystemTime::now().duration_since(UNIX_EPOCH) {
                 Ok(n) => n.as_secs().to_string(),
@@ -128,28 +145,15 @@ fn main() -> Result<(), Box<dyn Error>> {
                 }
             };
 
-            if args.is_present("extension") {
-                let ext = args.value_of("extension").unwrap().trim();
-                match ext {
-                    "jpeg" | "jpg" => {
-                        extension = backend::EncodingFormat::Jpg;
-                        log::debug!("Using custom extension: {:#?}", extension);
-                    }
-                    "png" => {}
-                    _ => {
-                        log::error!("Invalid extension provided.\nValid extensions:\n1) jpeg\n2) jpg\n3) png");
-                        exit(1);
-                    }
-                }
+            time + match extension {
+                backend::EncodingFormat::Png => "-wayshot.png",
+                backend::EncodingFormat::Jpg => "-wayshot.jpg",
+                backend::EncodingFormat::Ppm => "-wayshot.ppm",
             }
-            path = time
-                + match extension {
-                    backend::EncodingFormat::Png => "-wayshot.png",
-                    backend::EncodingFormat::Jpg => "-wayshot.jpg",
-                };
-        }
+        };
 
         backend::write_to_file(File::create(path)?, extension, frame_copy)?;
     }
+
     Ok(())
 }