New option --image-colors to specify colors used in image backends

o2sh · Oct 17, 2020 · 83459a8 · 83459a8
1 parent 7ccc949
commit 83459a8
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 5 deletions.
diff --git a/src/onefetch/cli.rs b/src/onefetch/cli.rs
@@ -16,6 +16,7 @@ pub struct Cli {
     pub no_bold: bool,
     pub image: Option<DynamicImage>,
     pub image_backend: Option<Box<dyn image_backends::ImageBackend>>,
+    pub image_colors: usize,
     pub no_merges: bool,
     pub no_color_blocks: bool,
     pub number_of_authors: usize,
@@ -115,6 +116,15 @@ impl Cli {
                 .possible_values(&possible_backends)
                 .help("Which image BACKEND to use."),
         )
+        .arg(
+            Arg::with_name("image-colors")
+                .long("image-colors")
+                .value_name("NUM")
+                .takes_value(true)
+                .max_values(1)
+                .default_value("16")
+                .help("NUM of colors (16-256) used in image backend."),
+        )
         .arg(
             Arg::with_name("no-merge-commits")
                 .long("no-merge-commits")
@@ -173,6 +183,17 @@ impl Cli {
             None
         };
 
+        let image_colors = if let Some(value) = matches.value_of("image-colors") {
+            let colors = usize::from_str(value).unwrap();
+            if 16 <= colors && colors <= 256 {
+                colors
+            } else {
+                16
+            }
+        } else {
+            16
+        };
+
         let path = String::from(matches.value_of("input").unwrap());
         let ascii_language = if let Some(ascii_language) = matches.value_of("ascii-language") {
             Language::from_str(&ascii_language.to_lowercase()).unwrap()
@@ -208,6 +229,7 @@ impl Cli {
             no_bold,
             image,
             image_backend,
+            image_colors,
             no_merges,
             no_color_blocks,
             number_of_authors,

diff --git a/src/onefetch/image_backends/kitty.rs b/src/onefetch/image_backends/kitty.rs
@@ -77,7 +77,7 @@ impl KittyBackend {
 }
 
 impl super::ImageBackend for KittyBackend {
-    fn add_image(&self, lines: Vec<String>, image: &DynamicImage) -> String {
+    fn add_image(&self, lines: Vec<String>, image: &DynamicImage, _colors: usize) -> String {
         let tty_size = unsafe {
             let tty_size: winsize = std::mem::zeroed();
             ioctl(STDOUT_FILENO, TIOCGWINSZ, &tty_size);

diff --git a/src/onefetch/image_backends/mod.rs b/src/onefetch/image_backends/mod.rs
@@ -6,7 +6,7 @@ pub mod kitty;
 pub mod sixel;
 
 pub trait ImageBackend {
-    fn add_image(&self, lines: Vec<String>, image: &DynamicImage) -> String;
+    fn add_image(&self, lines: Vec<String>, image: &DynamicImage, colors: usize) -> String;
 }
 
 #[cfg(not(windows))]

diff --git a/src/onefetch/image_backends/sixel.rs b/src/onefetch/image_backends/sixel.rs
@@ -71,7 +71,7 @@ impl SixelBackend {
 }
 
 impl super::ImageBackend for SixelBackend {
-    fn add_image(&self, lines: Vec<String>, image: &DynamicImage) -> String {
+    fn add_image(&self, lines: Vec<String>, image: &DynamicImage, colors: usize) -> String {
         let tty_size = unsafe {
             let tty_size: winsize = std::mem::zeroed();
             ioctl(STDOUT_FILENO, TIOCGWINSZ, &tty_size);
@@ -95,7 +95,7 @@ impl super::ImageBackend for SixelBackend {
         // reduce the amount of colors using dithering
         colorops::dither(
             &mut rgba_image,
-            &NeuQuant::new(10, 128, flat_samples.image_slice().unwrap()),
+            &NeuQuant::new(10, colors, flat_samples.image_slice().unwrap()),
         );
         let rgb_image = ImageBuffer::from_fn(rgba_image.width(), rgba_image.height(), |x, y| {
             let rgba_pixel = rgba_image.get_pixel(x, y);

diff --git a/src/onefetch/info.rs b/src/onefetch/info.rs
@@ -263,7 +263,8 @@ impl std::fmt::Display for Info {
                     "{}",
                     image_backend.add_image(
                         info_lines.map(|s| format!("{}{}", center_pad, s)).collect(),
-                        custom_image
+                        custom_image,
+                        self.config.image_colors
                     )
                 )?;
             } else {