From bbd1d47ba5798b994e723d05ccf410525b89a6e6 Mon Sep 17 00:00:00 2001
From: Dmitriy Pleshevskiy <dmitriy@ideascup.me>
Date: Sat, 12 Nov 2022 16:25:35 +0300
Subject: [PATCH] cli: add color option

---
 Cargo.lock        |  1 +
 Cargo.toml        |  3 ++-
 src/cli.rs        | 61 ++++++++++++++++++++++++++++++++++++++++++++++-
 src/cli/switch.rs | 28 ++++++++++++++--------
 4 files changed, 81 insertions(+), 12 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index f8a41c1..a0dec75 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -196,6 +196,7 @@ checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
 name = "vnetod"
 version = "0.3.3"
 dependencies = [
+ "atty",
  "clap",
  "termcolor",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 9ea72fe..046604e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -13,7 +13,8 @@ categories = ["command-line-interface", "config", "development-tools"]
 
 [dependencies]
 clap = { version = "3.2.15", default-features = false, features = ["std", "env", "derive"] }
+atty = { version = "0.2.14", optional = true }
 termcolor = { version = "1.1.3", optional = true }
 
 [features]
-color = ["clap/color", "dep:termcolor"]
+color = ["clap/color", "dep:atty", "dep:termcolor"]
diff --git a/src/cli.rs b/src/cli.rs
index e0d0b2a..ba2bf49 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -20,7 +20,7 @@ use std::path::PathBuf;
 
 use clap::Parser;
 
-#[derive(Parser)]
+#[derive(Parser, Debug)]
 #[clap(
     author,
     version,
@@ -62,4 +62,63 @@ pub struct Args {
         help = "Environment varible sections that will be enabled"
     )]
     pub sections: Vec<String>,
+
+    #[cfg(feature = "color")]
+    #[clap(
+        long,
+        value_enum,
+        default_value = "auto",
+        long_help = "
+This flag controls when to use colors. The default setting is 'auto', which
+means vnetod will try to guess when to use colors. For example, if vnetod is
+printing to a terminal, then it will use colors, but if it is redirected to a
+file or a pipe, then it will suppress color output. vnetod will suppress color
+output in some other circumstances as well. For example, if the TERM
+environment variable is not set or set to 'dumb', then vnetod will not use
+colors.
+
+The possible values for this flag are:
+
+never    Colors will never be used.
+auto     The default. vnetod tries to be smart.
+always   Colors will always be used regardless of where output is sent.
+ansi     Like 'always', but emits ANSI escapes (even in a Windows console).
+
+"
+    )]
+    pub color: ColorVariant,
+}
+
+#[cfg(feature = "color")]
+#[derive(clap::ValueEnum, Clone, Debug)]
+pub enum ColorVariant {
+    Auto,
+    Always,
+    Ansi,
+    Never,
+}
+
+#[cfg(feature = "color")]
+impl From<ColorVariant> for termcolor::ColorChoice {
+    fn from(col: ColorVariant) -> Self {
+        match col {
+            ColorVariant::Never => Self::Never,
+            ColorVariant::Always => Self::Always,
+            ColorVariant::Ansi => Self::AlwaysAnsi,
+            ColorVariant::Auto => {
+                if atty::is(atty::Stream::Stdout) {
+                    // Otherwise let's `termcolor` decide by inspecting the environment. From the [doc]:
+                    // - If `NO_COLOR` is set to any value, then colors will be suppressed.
+                    // - If `TERM` is set to dumb, then colors will be suppressed.
+                    // - In non-Windows environments, if `TERM` is not set, then colors will be suppressed.
+                    //
+                    // [doc]: https://github.com/BurntSushi/termcolor#automatic-color-selection
+                    Self::Auto
+                } else {
+                    // Colors should be deactivated if the terminal is not a tty.
+                    Self::Never
+                }
+            }
+        }
+    }
 }
diff --git a/src/cli/switch.rs b/src/cli/switch.rs
index 27e8a13..b20f699 100644
--- a/src/cli/switch.rs
+++ b/src/cli/switch.rs
@@ -15,7 +15,7 @@
 //!
 
 #[cfg(feature = "color")]
-use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
+use termcolor::{Color, ColorSpec, StandardStream, WriteColor};
 
 use crate::{cli::Args, domain};
 use std::fs::File;
@@ -47,11 +47,24 @@ pub fn execute(args: &Args) -> Result<(), Error> {
         })
         .transpose()?;
 
+    #[cfg(feature = "color")]
+    let color = args.color.clone();
+
+    println!();
+
     let new_content = domain::switch::execute(domain::switch::Request {
-        content: &content,
+        content: &content.trim(),
         sections: &args.sections,
-        on_line: Some(Box::new(print_line)),
+        on_line: Some(Box::new(move |line| {
+            #[cfg(feature = "color")]
+            print_line(line, color.clone());
+            #[cfg(not(feature = "color"))]
+            print!("{}", line)
+        })),
     });
+
+    println!();
+
     if let Some(mut fs_writer) = fs_writer {
         fs_writer
             .write_all(new_content.as_bytes())
@@ -62,8 +75,8 @@ pub fn execute(args: &Args) -> Result<(), Error> {
 }
 
 #[cfg(feature = "color")]
-fn print_line(line: &String) {
-    let mut stdout = StandardStream::stdout(ColorChoice::Always);
+fn print_line(line: &String, color: crate::cli::ColorVariant) {
+    let mut stdout = StandardStream::stdout(color.into());
     let color = line
         .starts_with("###")
         .then_some(Color::Yellow)
@@ -72,8 +85,3 @@ fn print_line(line: &String) {
     write!(&mut stdout, "{}", line).unwrap();
     stdout.reset().ok();
 }
-
-#[cfg(not(feature = "color"))]
-fn print_line(line: &String) {
-    print!("{}", line)
-}