Merge pull request #113 from CephalonRho/display-image

Add support for displaying a custom image instead of ascii art

Author: o2sh

Diff for: Cargo.toml

@@ -17,6 +17,11 @@ bytecount = "0.5.1"
 clap = "2.33.0"
 strum = "0.16.0"
 strum_macros = "0.16.0"
+image = "0.22.3"
 
 [target.'cfg(windows)'.dependencies]
 ansi_term = "0.12"
+
+[target.'cfg(target_os = "linux")'.dependencies]
+libc = "0.2.65"
+base64 = "0.11.0"

Diff for: src/error.rs

@@ -12,6 +12,8 @@ pub enum Error {
     NotGitRepo,
     /// Error while getting branch info
     ReferenceInfoError,
+    /// Image probably doesn't exist or has wrong format
+    ImageLoadError,
 }
 
 impl std::fmt::Debug for Error {
@@ -23,6 +25,7 @@ impl std::fmt::Debug for Error {
             Error::ReadDirectory => "Could not read directory",
             Error::NotGitRepo => "This is not a Git Repo",
             Error::ReferenceInfoError => "Error while retrieving reference information",
+            Error::ImageLoadError => "Could not load the specified image",
         };
         write!(f, "{}", content)
     }

Diff for: src/image_backends/kitty.rs

@@ -0,0 +1,139 @@
+use image::{imageops::FilterType, DynamicImage, GenericImageView};
+use libc::{
+    ioctl, tcgetattr, tcsetattr, termios, winsize, ECHO, ICANON, STDIN_FILENO, STDOUT_FILENO,
+    TCSANOW, TIOCGWINSZ,
+};
+use std::io::Read;
+use std::sync::mpsc::{self, TryRecvError};
+use std::time::Duration;
+
+pub struct KittyBackend {}
+
+impl KittyBackend {
+    pub fn new() -> Self {
+        Self {}
+    }
+
+    pub fn supported() -> bool {
+        // save terminal attributes and disable canonical input processing mode
+        let old_attributes = unsafe {
+            let mut old_attributes: termios = std::mem::zeroed();
+            tcgetattr(STDIN_FILENO, &mut old_attributes);
+
+            let mut new_attributes = old_attributes.clone();
+            new_attributes.c_lflag &= !ICANON;
+            new_attributes.c_lflag &= !ECHO;
+            tcsetattr(STDIN_FILENO, TCSANOW, &new_attributes);
+            old_attributes
+        };
+
+        // generate red rgba test image
+        let mut test_image = Vec::<u8>::with_capacity(32 * 32 * 4);
+        test_image.extend(
+            std::iter::repeat([255, 0, 0, 255].iter())
+                .take(32 * 32)
+                .flatten(),
+        );
+
+        // print the test image with the action set to query
+        println!(
+            "\x1B_Gi=1,f=32,s=32,v=32,a=q;{}\x1B\\",
+            base64::encode(&test_image)
+        );
+
+        // a new thread is required to avoid blocking the main thread if the terminal doesn't respond
+        let (sender, receiver) = mpsc::channel::<()>();
+        let (stop_sender, stop_receiver) = mpsc::channel::<()>();
+        std::thread::spawn(move || {
+            let mut buf = Vec::<u8>::new();
+            let allowed_bytes = [0x1B, '_' as u8, 'G' as u8, '\\' as u8];
+            for byte in std::io::stdin().lock().bytes() {
+                let byte = byte.unwrap();
+                if allowed_bytes.contains(&byte) {
+                    buf.push(byte);
+                }
+                if buf.starts_with(&[0x1B, '_' as u8, 'G' as u8])
+                    && buf.ends_with(&[0x1B, '\\' as u8])
+                {
+                    sender.send(()).unwrap();
+                    return;
+                }
+                match stop_receiver.try_recv() {
+                    Err(TryRecvError::Empty) => {}
+                    _ => return,
+                }
+            }
+        });
+        if let Ok(_) = receiver.recv_timeout(Duration::from_millis(50)) {
+            unsafe {
+                tcsetattr(STDIN_FILENO, TCSANOW, &old_attributes);
+            }
+            true
+        } else {
+            unsafe {
+                tcsetattr(STDIN_FILENO, TCSANOW, &old_attributes);
+            }
+            stop_sender.send(()).ok();
+            false
+        }
+    }
+}
+
+impl super::ImageBackend for KittyBackend {
+    fn add_image(&self, lines: Vec<String>, image: &DynamicImage) -> String {
+        let tty_size = unsafe {
+            let tty_size: winsize = std::mem::zeroed();
+            ioctl(STDOUT_FILENO, TIOCGWINSZ, &tty_size);
+            tty_size
+        };
+        let width_ratio = tty_size.ws_col as f64 / tty_size.ws_xpixel as f64;
+        let height_ratio = tty_size.ws_row as f64 / tty_size.ws_ypixel as f64;
+
+        // resize image to fit the text height with the Lanczos3 algorithm
+        let image = image.resize(
+            u32::max_value(),
+            (lines.len() as f64 / height_ratio) as u32,
+            FilterType::Lanczos3,
+        );
+        let _image_columns = width_ratio * image.width() as f64;
+        let image_rows = height_ratio * image.height() as f64;
+
+        // convert the image to rgba samples
+        let rgba_image = image.to_rgba();
+        let flat_samples = rgba_image.as_flat_samples();
+        let raw_image = flat_samples
+            .image_slice()
+            .expect("Conversion from image to rgba samples failed");
+        assert_eq!(
+            image.width() as usize * image.height() as usize * 4,
+            raw_image.len()
+        );
+        
+        let encoded_image = base64::encode(&raw_image); // image data is base64 encoded
+        let mut image_data = Vec::<u8>::new();
+        for chunk in encoded_image.as_bytes().chunks(4096) {
+            // send a 4096 byte chunk of base64 encoded rgba image data
+            image_data.extend(
+                format!(
+                    "\x1B_Gf=32,s={},v={},m=1,a=T;",
+                    image.width(),
+                    image.height()
+                )
+                .as_bytes(),
+            );
+            image_data.extend(chunk);
+            image_data.extend("\x1B\\".as_bytes());
+        }
+        image_data.extend("\x1B_Gm=0;\x1B\\".as_bytes()); // write empty last chunk
+        image_data.extend(format!("\x1B[{}A", image_rows as u32 - 1).as_bytes()); // move cursor to start of image
+        let mut i = 0;
+        for line in &lines {
+            image_data.extend(format!("\x1B[s{}\x1B[u\x1B[1B", line).as_bytes());
+            i += 1;
+        }
+        image_data
+            .extend(format!("\n\x1B[{}B", lines.len().max(image_rows as usize) - i).as_bytes()); // move cursor to end of image
+
+        String::from_utf8(image_data).unwrap()
+    }
+}

Diff for: src/image_backends/mod.rs

@@ -0,0 +1,22 @@
+use image::DynamicImage;
+
+#[cfg(target_os = "linux")]
+mod kitty;
+
+pub trait ImageBackend {
+    fn add_image(&self, lines: Vec<String>, image: &DynamicImage) -> String;
+}
+
+#[cfg(target_os = "linux")]
+pub fn get_best_backend() -> Option<Box<dyn ImageBackend>> {
+    if kitty::KittyBackend::supported() {
+        Some(Box::new(kitty::KittyBackend::new()))
+    } else {
+        None
+    }
+}
+
+#[cfg(not(target_os = "linux"))]
+pub fn get_best_backend() -> Option<Box<dyn ImageBackend>> {
+    None
+}

Diff for: src/info.rs

@@ -7,9 +7,11 @@ use std::str::FromStr;
 use colored::{Color, Colorize, ColoredString};
 use git2::Repository;
 use license::License;
+use image::DynamicImage;
 
 use crate::language::Language;
 use crate::{AsciiArt, CommitInfo, Configuration, Error, InfoFieldOn};
+use crate::image_backends;
 
 type Result<T> = std::result::Result<T, crate::Error>;
 
@@ -33,6 +35,7 @@ pub struct Info {
     custom_colors: Vec<String>,
     disable_fields: InfoFieldOn,
     bold_enabled: bool,
+    custom_image: Option<DynamicImage>,
 }
 
 impl std::fmt::Display for Info {
@@ -177,27 +180,35 @@ impl std::fmt::Display for Info {
             "   ".on_bright_white(),
         )?;
 
-        let mut logo_lines = AsciiArt::new(self.get_ascii(), self.colors(), self.bold_enabled);
+        let center_pad = "   ";
         let mut info_lines = buf.lines();
 
-        let center_pad = "   ";
-        loop {
-            match (logo_lines.next(), info_lines.next()) {
-                (Some(logo_line), Some(info_line)) => {
-                    writeln!(f, "{}{}{:^}", logo_line, center_pad, info_line)?
-                }
-                (Some(logo_line), None) => writeln!(f, "{}", logo_line)?,
-                (None, Some(info_line)) => writeln!(
-                    f,
-                    "{:<width$}{}{:^}",
-                    "",
-                    center_pad,
-                    info_line,
-                    width = logo_lines.width()
-                )?,
-                (None, None) => {
-                    writeln!(f, "\n")?;
-                    break;
+        if let Some(custom_image) = &self.custom_image {
+            if let Some(backend) = image_backends::get_best_backend() {
+                writeln!(f, "{}", backend.add_image(info_lines.map(|s| format!("{}{}", center_pad, s)).collect(), custom_image))?;
+            } else {
+                panic!("No image backend found")
+            }
+        } else {
+            let mut logo_lines = AsciiArt::new(self.get_ascii(), self.colors(), self.bold_enabled);
+            loop {
+                match (logo_lines.next(), info_lines.next()) {
+                    (Some(logo_line), Some(info_line)) => {
+                        writeln!(f, "{}{}{:^}", logo_line, center_pad, info_line)?
+                    }
+                    (Some(logo_line), None) => writeln!(f, "{}", logo_line)?,
+                    (None, Some(info_line)) => writeln!(
+                        f,
+                        "{:<width$}{}{:^}",
+                        "",
+                        center_pad,
+                        info_line,
+                        width = logo_lines.width()
+                    )?,
+                    (None, None) => {
+                        writeln!(f, "\n")?;
+                        break;
+                    }
                 }
             }
         }
@@ -213,6 +224,7 @@ impl Info {
         colors: Vec<String>,
         disabled: InfoFieldOn,
         bold_flag: bool,
+        custom_image: Option<DynamicImage>
     ) -> Result<Info> {
         let authors = Info::get_authors(&dir, 3);
         let (git_v, git_user) = Info::get_git_info(&dir);
@@ -247,6 +259,7 @@ impl Info {
             custom_colors: colors,
             disable_fields: disabled,
             bold_enabled: bold_flag,
+            custom_image,
         })
     }
 

Diff for: src/main.rs

@@ -8,10 +8,14 @@ extern crate clap;
 extern crate strum;
 #[macro_use]
 extern crate strum_macros;
+extern crate image;
 
 #[cfg(target = "windows")]
 extern crate ansi_term;
 
+#[cfg(target_os = "linux")]
+extern crate libc;
+
 use clap::{App, Arg};
 use colored::*;
 use std::{
@@ -25,6 +29,7 @@ use strum::{EnumCount, IntoEnumIterator};
 mod ascii_art;
 mod commit_info;
 mod error;
+mod image_backends;
 mod info;
 mod language;
 
@@ -165,6 +170,12 @@ Possible values: [{0}{1}{2}{3}{4}{5}{6}{7}{8}{9}{10}{11}{12}{13}{14}{15}]",
                 .short("l")
                 .long("languages")
                 .help("Prints out supported languages"),
+        ).arg(
+            Arg::with_name("image")
+                .short("i")
+                .long("image")
+                .takes_value(true)
+                .help("Sets a custom image to use instead of the ascii logo"),
         )
         .get_matches();
 
@@ -220,7 +231,13 @@ Possible values: [{0}{1}{2}{3}{4}{5}{6}{7}{8}{9}{10}{11}{12}{13}{14}{15}]",
 
     let bold_flag = !matches.is_present("no-bold");
 
-    let info = Info::new(&dir, custom_logo, custom_colors, disable_fields, bold_flag)?;
+    let custom_image = if let Some(image_path) = matches.value_of("image") {
+        Some(image::open(image_path).map_err(|_| Error::ImageLoadError)?)
+    } else {
+        None
+    };
+
+    let info = Info::new(&dir, custom_logo, custom_colors, disable_fields, bold_flag, custom_image)?;
 
     print!("{}", info);
     Ok(())