Add dmi_create_png function (tgstation#21)

To go around the painting problem and maybe other things in the future.
ParadiseSS13 · Feb 6, 2020 · ab4a02a · ab4a02a
1 parent f22afdd
commit ab4a02a
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/dmi.rs b/src/dmi.rs
@@ -1,13 +1,16 @@
 use std::fs::File;
-
 use png::{Decoder, Encoder, HasParameters, OutputInfo};
 
-use error::Result;
+use error::{Result, Error};
 
 byond_fn! { dmi_strip_metadata(path) {
     strip_metadata(path).err()
 } }
 
+byond_fn! { dmi_create_png(path, width, height, data) {
+    create_png(path, width, height , data).err()
+} }
+
 fn strip_metadata(path: &str) -> Result<()> {
     let (info, image) = read_png(path)?;
     Ok(write_png(path, info, image)?)
@@ -28,3 +31,26 @@ fn write_png(path: &str, info: OutputInfo, image: Vec<u8>) -> Result<()> {
     let mut writer = encoder.write_header()?;
     Ok(writer.write_image_data(&image)?)
 }
+
+fn create_png(path: &str, width: &str, height: &str, data: &str) -> Result<()> {
+    let width = u32::from_str_radix(width, 10)?;
+    let height = u32::from_str_radix(height, 10)?;
+
+    let bytes = data.as_bytes();
+    if bytes.len() % 7 != 0 {
+        return Err(Error::InvalidPngDataError);
+    }
+
+    let mut result: Vec<u8> = Vec::new();
+    for pixel in bytes.chunks_exact(7) {
+        for channel in pixel[1..].chunks_exact(2) {
+            result.push(u8::from_str_radix(std::str::from_utf8(channel)?, 16)?);
+        }
+    }
+
+    let mut encoder = Encoder::new(File::create(path)?, width, height);
+    encoder.set(png::ColorType::RGB);
+    encoder.set(png::BitDepth::Eight);
+    let mut writer = encoder.write_header()?;
+    Ok(writer.write_image_data(&result)?)
+}
diff --git a/src/error.rs b/src/error.rs
@@ -1,6 +1,7 @@
 use std::io;
 use std::result;
 use std::str::Utf8Error;
+use std::num::ParseIntError;
 
 #[cfg(feature="png")]
 use png::{DecodingError, EncodingError};
@@ -25,6 +26,12 @@ pub enum Error {
     #[cfg(feature="png")]
     #[fail(display = "{}", _0)]
     ImageEncoding(#[cause] EncodingError),
+    #[cfg(feature="png")]
+    #[fail(display = "{}", _0)]
+    ParseIntError(#[cause] ParseIntError),
+    #[cfg(feature="png")]
+    #[fail(display = "Invalid png data.")]
+    InvalidPngDataError,
 }
 
 impl From<io::Error> for Error {
@@ -53,6 +60,13 @@ impl From<EncodingError> for Error {
     }
 }
 
+#[cfg(feature="png")]
+impl From<ParseIntError> for Error {
+    fn from(error: ParseIntError) -> Error {
+        Error::ParseIntError(error)
+    }
+}
+
 impl From<Error> for String {
     fn from(error: Error) -> String {
         error.to_string()