@@ -0,0 +1,120 @@ |
| 1 | +//! PNG encoding for garshot. |
| 2 | + |
| 3 | +use std::fs::File; |
| 4 | +use std::io::BufWriter; |
| 5 | +use std::path::Path; |
| 6 | + |
| 7 | +use png::{BitDepth, ColorType, Compression, Encoder, FilterType}; |
| 8 | + |
| 9 | +use crate::error::{GarshotError, Result}; |
| 10 | + |
| 11 | +/// Encode RGBA image data to a PNG file. |
| 12 | +/// |
| 13 | +/// Uses fast compression by default for better screenshot responsiveness. |
| 14 | +/// |
| 15 | +/// # Arguments |
| 16 | +/// * `data` - RGBA pixel data (4 bytes per pixel) |
| 17 | +/// * `width` - Image width in pixels |
| 18 | +/// * `height` - Image height in pixels |
| 19 | +/// * `path` - Output file path |
| 20 | +pub fn encode_png(data: &[u8], width: u32, height: u32, path: &Path) -> Result<()> { |
| 21 | + let expected_size = (width * height * 4) as usize; |
| 22 | + if data.len() != expected_size { |
| 23 | + return Err(GarshotError::EncodeError(format!( |
| 24 | + "Data size {} does not match expected size {} for {}x{} RGBA image", |
| 25 | + data.len(), |
| 26 | + expected_size, |
| 27 | + width, |
| 28 | + height |
| 29 | + ))); |
| 30 | + } |
| 31 | + |
| 32 | + let file = File::create(path)?; |
| 33 | + let writer = BufWriter::new(file); |
| 34 | + |
| 35 | + let mut encoder = Encoder::new(writer, width, height); |
| 36 | + encoder.set_color(ColorType::Rgba); |
| 37 | + encoder.set_depth(BitDepth::Eight); |
| 38 | + |
| 39 | + // Use fast compression for responsiveness |
| 40 | + // Screenshots are typically viewed immediately, so speed > size |
| 41 | + encoder.set_compression(Compression::Fast); |
| 42 | + encoder.set_filter(FilterType::Sub); // Good for screenshots |
| 43 | + |
| 44 | + let mut png_writer = encoder |
| 45 | + .write_header() |
| 46 | + .map_err(|e| GarshotError::EncodeError(e.to_string()))?; |
| 47 | + |
| 48 | + png_writer |
| 49 | + .write_image_data(data) |
| 50 | + .map_err(|e| GarshotError::EncodeError(e.to_string()))?; |
| 51 | + |
| 52 | + tracing::debug!("Encoded PNG {}x{} to {}", width, height, path.display()); |
| 53 | + |
| 54 | + Ok(()) |
| 55 | +} |
| 56 | + |
| 57 | +/// Encode RGBA image data to PNG bytes in memory. |
| 58 | +/// |
| 59 | +/// Useful for stdout output or clipboard operations. |
| 60 | +pub fn encode_png_to_vec(data: &[u8], width: u32, height: u32) -> Result<Vec<u8>> { |
| 61 | + let expected_size = (width * height * 4) as usize; |
| 62 | + if data.len() != expected_size { |
| 63 | + return Err(GarshotError::EncodeError(format!( |
| 64 | + "Data size {} does not match expected size {} for {}x{} RGBA image", |
| 65 | + data.len(), |
| 66 | + expected_size, |
| 67 | + width, |
| 68 | + height |
| 69 | + ))); |
| 70 | + } |
| 71 | + |
| 72 | + let mut buffer = Vec::new(); |
| 73 | + |
| 74 | + { |
| 75 | + let mut encoder = Encoder::new(&mut buffer, width, height); |
| 76 | + encoder.set_color(ColorType::Rgba); |
| 77 | + encoder.set_depth(BitDepth::Eight); |
| 78 | + encoder.set_compression(Compression::Fast); |
| 79 | + encoder.set_filter(FilterType::Sub); |
| 80 | + |
| 81 | + let mut png_writer = encoder |
| 82 | + .write_header() |
| 83 | + .map_err(|e| GarshotError::EncodeError(e.to_string()))?; |
| 84 | + |
| 85 | + png_writer |
| 86 | + .write_image_data(data) |
| 87 | + .map_err(|e| GarshotError::EncodeError(e.to_string()))?; |
| 88 | + } |
| 89 | + |
| 90 | + Ok(buffer) |
| 91 | +} |
| 92 | + |
| 93 | +#[cfg(test)] |
| 94 | +mod tests { |
| 95 | + use super::*; |
| 96 | + use std::io::Read; |
| 97 | + |
| 98 | + #[test] |
| 99 | + fn test_encode_png_to_vec() { |
| 100 | + // Create a 2x2 red image |
| 101 | + let data = vec![ |
| 102 | + 255, 0, 0, 255, // Red pixel |
| 103 | + 255, 0, 0, 255, // Red pixel |
| 104 | + 255, 0, 0, 255, // Red pixel |
| 105 | + 255, 0, 0, 255, // Red pixel |
| 106 | + ]; |
| 107 | + |
| 108 | + let png_data = encode_png_to_vec(&data, 2, 2).unwrap(); |
| 109 | + |
| 110 | + // Check PNG magic bytes |
| 111 | + assert_eq!(&png_data[0..8], &[137, 80, 78, 71, 13, 10, 26, 10]); |
| 112 | + } |
| 113 | + |
| 114 | + #[test] |
| 115 | + fn test_encode_png_invalid_size() { |
| 116 | + let data = vec![0; 10]; // Wrong size |
| 117 | + let result = encode_png_to_vec(&data, 2, 2); |
| 118 | + assert!(result.is_err()); |
| 119 | + } |
| 120 | +} |