Better handling for color in terminal (#381)

* better handling for color in terminal

* cleanup

* cleanup

* cargo fmt

* clippy + tests

* clean

11 files changed, 73 insertions, 22 deletions

diff --git a/README.md b/README.md
index d2dfc4d..effb292 100644
--- a/README.md
+++ b/README.md
@@ -76,9 +76,10 @@ Usage: dust -o si/b/kb/kib/mb/mib/gb/gib (si - prints sizes in powers of 1000. O
 Usage: dust -X ignore  (ignore all files and directories with the name 'ignore')
 Usage: dust -x (Only show directories on the same filesystem)
 Usage: dust -b (Do not show percentages or draw ASCII bars)
-Usage: dust -B (--bars-on-right - Percent bars moved to right side of screen])
+Usage: dust -B (--bars-on-right - Percent bars moved to right side of screen)
 Usage: dust -i (Do not show hidden files)
 Usage: dust -c (No colors [monochrome])
+Usage: dust -C (Force colors)
 Usage: dust -f (Count files instead of diskspace)
 Usage: dust -t (Group by filetype)
 Usage: dust -z 10M (min-size, Only include files larger than 10M)
diff --git a/completions/_dust b/completions/_dust
index 4b850d7..ecf5150 100644
--- a/completions/_dust
+++ b/completions/_dust
@@ -47,6 +47,8 @@ _dust() {
 '--reverse[Print tree upside down (biggest highest)]' \
 '-c[No colors will be printed (Useful for commands like\: watch)]' \
 '--no-colors[No colors will be printed (Useful for commands like\: watch)]' \
+'-C[Force colors print]' \
+'--force-colors[Force colors print]' \
 '-b[No percent bars or percentages will be displayed]' \
 '--no-percent-bars[No percent bars or percentages will be displayed]' \
 '-B[percent bars moved to right side of screen]' \
diff --git a/completions/_dust.ps1 b/completions/_dust.ps1
index ddfb26f..484d20e 100644
--- a/completions/_dust.ps1
+++ b/completions/_dust.ps1
@@ -53,6 +53,8 @@ Register-ArgumentCompleter -Native -CommandName 'dust' -ScriptBlock {
             [CompletionResult]::new('--reverse', 'reverse', [CompletionResultType]::ParameterName, 'Print tree upside down (biggest highest)')
             [CompletionResult]::new('-c', 'c', [CompletionResultType]::ParameterName, 'No colors will be printed (Useful for commands like: watch)')
             [CompletionResult]::new('--no-colors', 'no-colors', [CompletionResultType]::ParameterName, 'No colors will be printed (Useful for commands like: watch)')
+            [CompletionResult]::new('-C', 'C ', [CompletionResultType]::ParameterName, 'Force colors print')
+            [CompletionResult]::new('--force-colors', 'force-colors', [CompletionResultType]::ParameterName, 'Force colors print')
             [CompletionResult]::new('-b', 'b', [CompletionResultType]::ParameterName, 'No percent bars or percentages will be displayed')
             [CompletionResult]::new('--no-percent-bars', 'no-percent-bars', [CompletionResultType]::ParameterName, 'No percent bars or percentages will be displayed')
             [CompletionResult]::new('-B', 'B ', [CompletionResultType]::ParameterName, 'percent bars moved to right side of screen')
diff --git a/completions/dust.bash b/completions/dust.bash
index 24c37e0..b332455 100644
--- a/completions/dust.bash
+++ b/completions/dust.bash
@@ -19,7 +19,7 @@ _dust() {
 
     case "${cmd}" in
         dust)
-            opts="-d -n -p -X -I -L -x -s -r -c -b -B -z -R -f -i -v -e -t -w -P -D -F -o -S -h -V --depth --number-of-lines --full-paths --ignore-directory --ignore-all-in-file --dereference-links --limit-filesystem --apparent-size --reverse --no-colors --no-percent-bars --bars-on-right --min-size --screen-reader --skip-total --filecount --ignore_hidden --invert-filter --filter --file_types --terminal_width --no-progress --only-dir --only-file --output-format --stack-size --help --version [params]..."
+            opts="-d -n -p -X -I -L -x -s -r -c -C -b -B -z -R -f -i -v -e -t -w -P -D -F -o -S -h -V --depth --number-of-lines --full-paths --ignore-directory --ignore-all-in-file --dereference-links --limit-filesystem --apparent-size --reverse --no-colors --force-colors --no-percent-bars --bars-on-right --min-size --screen-reader --skip-total --filecount --ignore_hidden --invert-filter --filter --file_types --terminal_width --no-progress --only-dir --only-file --output-format --stack-size --help --version [params]..."
             if [[ ${cur} == -* || ${COMP_CWORD} -eq 1 ]] ; then
                 COMPREPLY=( $(compgen -W "${opts}" -- "${cur}") )
                 return 0
diff --git a/completions/dust.elv b/completions/dust.elv
index a001ccb..921c4e9 100644
--- a/completions/dust.elv
+++ b/completions/dust.elv
@@ -50,6 +50,8 @@ set edit:completion:arg-completer[dust] = {|@words|
             cand --reverse 'Print tree upside down (biggest highest)'
             cand -c 'No colors will be printed (Useful for commands like: watch)'
             cand --no-colors 'No colors will be printed (Useful for commands like: watch)'
+            cand -C 'Force colors print'
+            cand --force-colors 'Force colors print'
             cand -b 'No percent bars or percentages will be displayed'
             cand --no-percent-bars 'No percent bars or percentages will be displayed'
             cand -B 'percent bars moved to right side of screen'
diff --git a/completions/dust.fish b/completions/dust.fish
index 5dbd22d..4b09292 100644
--- a/completions/dust.fish
+++ b/completions/dust.fish
@@ -14,6 +14,7 @@ complete -c dust -s x -l limit-filesystem -d 'Only count the files and directori
 complete -c dust -s s -l apparent-size -d 'Use file length instead of blocks'
 complete -c dust -s r -l reverse -d 'Print tree upside down (biggest highest)'
 complete -c dust -s c -l no-colors -d 'No colors will be printed (Useful for commands like: watch)'
+complete -c dust -s C -l force-colors -d 'Force colors print'
 complete -c dust -s b -l no-percent-bars -d 'No percent bars or percentages will be displayed'
 complete -c dust -s B -l bars-on-right -d 'percent bars moved to right side of screen'
 complete -c dust -s R -l screen-reader -d 'For screen readers. Removes bars. Adds new column: depth level (May want to use -p too for full path)'
diff --git a/man-page/dust.1 b/man-page/dust.1
index e032029..5296cf4 100644
--- a/man-page/dust.1
+++ b/man-page/dust.1
@@ -4,7 +4,7 @@
 .SH NAME
 Dust \- Like du but more intuitive
 .SH SYNOPSIS
-\fBdust\fR [\fB\-d\fR|\fB\-\-depth\fR] [\fB\-n\fR|\fB\-\-number\-of\-lines\fR] [\fB\-p\fR|\fB\-\-full\-paths\fR] [\fB\-X\fR|\fB\-\-ignore\-directory\fR] [\fB\-I\fR|\fB\-\-ignore\-all\-in\-file\fR] [\fB\-L\fR|\fB\-\-dereference\-links\fR] [\fB\-x\fR|\fB\-\-limit\-filesystem\fR] [\fB\-s\fR|\fB\-\-apparent\-size\fR] [\fB\-r\fR|\fB\-\-reverse\fR] [\fB\-c\fR|\fB\-\-no\-colors\fR] [\fB\-b\fR|\fB\-\-no\-percent\-bars\fR] [\fB\-B\fR|\fB\-\-bars\-on\-right\fR] [\fB\-z\fR|\fB\-\-min\-size\fR] [\fB\-R\fR|\fB\-\-screen\-reader\fR] [\fB\-\-skip\-total\fR] [\fB\-f\fR|\fB\-\-filecount\fR] [\fB\-i\fR|\fB\-\-ignore_hidden\fR] [\fB\-v\fR|\fB\-\-invert\-filter\fR] [\fB\-e\fR|\fB\-\-filter\fR] [\fB\-t\fR|\fB\-\-file_types\fR] [\fB\-w\fR|\fB\-\-terminal_width\fR] [\fB\-P\fR|\fB\-\-no\-progress\fR] [\fB\-D\fR|\fB\-\-only\-dir\fR] [\fB\-F\fR|\fB\-\-only\-file\fR] [\fB\-o\fR|\fB\-\-output\-format\fR] [\fB\-S\fR|\fB\-\-stack\-size\fR] [\fB\-h\fR|\fB\-\-help\fR] [\fB\-V\fR|\fB\-\-version\fR] [\fIparams\fR] 
+\fBdust\fR [\fB\-d\fR|\fB\-\-depth\fR] [\fB\-n\fR|\fB\-\-number\-of\-lines\fR] [\fB\-p\fR|\fB\-\-full\-paths\fR] [\fB\-X\fR|\fB\-\-ignore\-directory\fR] [\fB\-I\fR|\fB\-\-ignore\-all\-in\-file\fR] [\fB\-L\fR|\fB\-\-dereference\-links\fR] [\fB\-x\fR|\fB\-\-limit\-filesystem\fR] [\fB\-s\fR|\fB\-\-apparent\-size\fR] [\fB\-r\fR|\fB\-\-reverse\fR] [\fB\-c\fR|\fB\-\-no\-colors\fR] [\fB\-C\fR|\fB\-\-force\-colors\fR] [\fB\-b\fR|\fB\-\-no\-percent\-bars\fR] [\fB\-B\fR|\fB\-\-bars\-on\-right\fR] [\fB\-z\fR|\fB\-\-min\-size\fR] [\fB\-R\fR|\fB\-\-screen\-reader\fR] [\fB\-\-skip\-total\fR] [\fB\-f\fR|\fB\-\-filecount\fR] [\fB\-i\fR|\fB\-\-ignore_hidden\fR] [\fB\-v\fR|\fB\-\-invert\-filter\fR] [\fB\-e\fR|\fB\-\-filter\fR] [\fB\-t\fR|\fB\-\-file_types\fR] [\fB\-w\fR|\fB\-\-terminal_width\fR] [\fB\-P\fR|\fB\-\-no\-progress\fR] [\fB\-D\fR|\fB\-\-only\-dir\fR] [\fB\-F\fR|\fB\-\-only\-file\fR] [\fB\-o\fR|\fB\-\-output\-format\fR] [\fB\-S\fR|\fB\-\-stack\-size\fR] [\fB\-h\fR|\fB\-\-help\fR] [\fB\-V\fR|\fB\-\-version\fR] [\fIparams\fR] 
 .SH DESCRIPTION
 Like du but more intuitive
 .SH OPTIONS
@@ -39,6 +39,9 @@ Print tree upside down (biggest highest)
 \fB\-c\fR, \fB\-\-no\-colors\fR
 No colors will be printed (Useful for commands like: watch)
 .TP
+\fB\-C\fR, \fB\-\-force\-colors\fR
+Force colors print
+.TP
 \fB\-b\fR, \fB\-\-no\-percent\-bars\fR
 No percent bars or percentages will be displayed
 .TP
diff --git a/src/cli.rs b/src/cli.rs
index 67a88f7..1e2f0d0 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -81,6 +81,13 @@ pub fn build_cli() -> Command {
                 .help("No colors will be printed (Useful for commands like: watch)"),
         )
         .arg(
+            Arg::new("force_colors")
+                .short('C')
+                .long("force-colors")
+                .action(clap::ArgAction::SetTrue)
+                .help("Force colors print"),
+        )
+        .arg(
             Arg::new("no_bars")
                 .short('b')
                 .long("no-percent-bars")
diff --git a/src/config.rs b/src/config.rs
index 1346fff..e1cd75b 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -16,6 +16,7 @@ pub struct Config {
     pub display_apparent_size: Option<bool>,
     pub reverse: Option<bool>,
     pub no_colors: Option<bool>,
+    pub force_colors: Option<bool>,
     pub no_bars: Option<bool>,
     pub skip_total: Option<bool>,
     pub screen_reader: Option<bool>,
@@ -34,6 +35,9 @@ impl Config {
     pub fn get_no_colors(&self, options: &ArgMatches) -> bool {
         Some(true) == self.no_colors || options.get_flag("no_colors")
     }
+    pub fn get_force_colors(&self, options: &ArgMatches) -> bool {
+        Some(true) == self.force_colors || options.get_flag("force_colors")
+    }
     pub fn get_disable_progress(&self, options: &ArgMatches) -> bool {
         Some(true) == self.disable_progress
             || options.get_flag("disable_progress")
diff --git a/src/main.rs b/src/main.rs
index 605e7b0..d4fe562 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -19,6 +19,7 @@ use filter::AggregateData;
 use progress::PIndicator;
 use regex::Error;
 use std::collections::HashSet;
+use std::env;
 use std::fs::read_to_string;
 use std::panic;
 use std::process;
@@ -41,29 +42,38 @@ use utils::simplify_dir_names;
 static DEFAULT_NUMBER_OF_LINES: usize = 30;
 static DEFAULT_TERMINAL_WIDTH: usize = 80;
 
-fn init_color(no_color: bool) -> bool {
+fn should_init_color(no_color: bool, force_color: bool) -> bool {
+    if force_color {
+        return true;
+    }
+    if no_color {
+        return false;
+    }
+    // check if NO_COLOR is set
+    // https://no-color.org/
+    if env::var_os("NO_COLOR").is_some() {
+        return false;
+    }
+    if terminal_size().is_none() {
+        // we are not in a terminal, color may not be needed
+        return false;
+    }
+    // we are in a terminal
     #[cfg(windows)]
     {
-        // If no color is already set do not print a warning message
-        if no_color {
-            true
-        } else {
-            // Required for windows 10
-            // Fails to resolve for windows 8 so disable color
-            match ansi_term::enable_ansi_support() {
-                Ok(_) => no_color,
-                Err(_) => {
-                    eprintln!(
-                    "This version of Windows does not support ANSI colors, setting no_color flag"
-                );
-                    true
-                }
+        // Required for windows 10
+        // Fails to resolve for windows 8 so disable color
+        return match ansi_term::enable_ansi_support() {
+            Ok(_) => true,
+            Err(_) => {
+                eprintln!("This version of Windows does not support ANSI colors");
+                false
             }
-        }
+        };
     }
     #[cfg(not(windows))]
     {
-        no_color
+        true
     }
 }
 
@@ -147,7 +157,10 @@ fn main() {
         }
     };
 
-    let no_colors = init_color(config.get_no_colors(&options));
+    let is_colors = should_init_color(
+        config.get_no_colors(&options),
+        config.get_force_colors(&options),
+    );
 
     let ignore_directories = match options.get_many::<String>("ignore_directory") {
         Some(values) => values
@@ -265,7 +278,7 @@ fn main() {
         let idd = InitialDisplayData {
             short_paths: !config.get_full_paths(&options),
             is_reversed: !config.get_reverse(&options),
-            colors_on: !no_colors,
+            colors_on: is_colors,
             by_filecount,
             is_screen_reader: config.get_screen_reader(&options),
             output_format,
diff --git a/tests/test_flags.rs b/tests/test_flags.rs
index 3251692..0a1f56e 100644
--- a/tests/test_flags.rs
+++ b/tests/test_flags.rs
@@ -232,3 +232,19 @@ pub fn test_show_files_by_invert_regex_match_multiple() {
     assert!(!output.contains("test_dir_unicode"));
     assert!(output.contains("many"));
 }
+
+#[test]
+pub fn test_no_color() {
+    let output = build_command(vec!["-c"]);
+    // Red is 31
+    assert!(!output.contains("\x1B[31m"));
+    assert!(!output.contains("\x1B[0m"));
+}
+
+#[test]
+pub fn test_force_color() {
+    let output = build_command(vec!["-C"]);
+    // Red is 31
+    assert!(output.contains("\x1B[31m"));
+    assert!(output.contains("\x1B[0m"));
+}