From f253e6eb18bc358a608f75129ae6ddfe8f511ee2 Mon Sep 17 00:00:00 2001
From: Zhaofeng Li <hello@zhaofeng.li>
Date: Sat, 4 Dec 2021 01:03:26 -0800
Subject: [PATCH] cli: Allow configuring output colorization

We now follow the <https://bixense.com/clicolors> standard.
---
 Cargo.lock  | 13 +++++++++++++
 Cargo.toml  |  1 +
 default.nix |  2 +-
 src/cli.rs  | 45 ++++++++++++++++++++++++++++++++++++++++++++-
 src/main.rs | 14 --------------
 5 files changed, 59 insertions(+), 16 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6feeb99..55632b3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -118,6 +118,18 @@ dependencies = [
  "vec_map",
 ]
 
+[[package]]
+name = "clicolors-control"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "90082ee5dcdd64dc4e9e0d37fbf3ee325419e39c0092191e0393df65518f741e"
+dependencies = [
+ "atty",
+ "lazy_static",
+ "libc",
+ "winapi",
+]
+
 [[package]]
 name = "colmena"
 version = "0.3.0-pre"
@@ -126,6 +138,7 @@ dependencies = [
  "async-trait",
  "atty",
  "clap",
+ "clicolors-control",
  "console 0.13.0",
  "const_format",
  "env_logger",
diff --git a/Cargo.toml b/Cargo.toml
index 69cbe7f..87a5ac5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,6 +11,7 @@ ansi-to-html = "0.1.0"
 async-trait = "0.1.42"
 atty = "0.2"
 clap = "2.33.3"
+clicolors-control = "1"
 console = "0.13.0"
 const_format = "0.2.22"
 env_logger = "0.8.2"
diff --git a/default.nix b/default.nix
index 57c60c8..3196446 100644
--- a/default.nix
+++ b/default.nix
@@ -20,7 +20,7 @@ in rustPlatform.buildRustPackage rec {
     src = lib.cleanSource ./.;
   };
 
-  cargoSha256 = "sha256-HGqecerb5LgnPhetqBYEmDKpJBkgzLS+iviVkDgVyGI=";
+  cargoSha256 = "sha256-YE+0jsdkzRaiUhPC71WBe/KcYERNy7nME26+ejrPRlc=";
 
   postInstall = lib.optionalString (stdenv.hostPlatform == stdenv.buildPlatform) ''
     mkdir completions
diff --git a/src/cli.rs b/src/cli.rs
index a5d6f0e..ba33bd3 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -1,7 +1,10 @@
 //! Global CLI Setup.
 
+use std::env;
+
 use clap::{App, AppSettings, Arg, ArgMatches, SubCommand};
 use const_format::concatcp;
+use env_logger::fmt::WriteStyle;
 use lazy_static::lazy_static;
 
 use crate::command;
@@ -104,7 +107,18 @@ pub fn build_cli(include_internal: bool) -> App<'static, 'static> {
             .help("Show debug information for Nix commands")
             .long_help("Passes --show-trace to Nix commands")
             .global(true)
-            .takes_value(false));
+            .takes_value(false))
+        .arg(Arg::with_name("color")
+            .long("color")
+            .help("When to colorize the output")
+            .long_help(r#"When to colorize the output. By default, Colmena enables colorized output when the terminal supports it.
+
+It's also possible to specify the preference using environment variables. See <https://bixense.com/clicolors>.
+"#)
+            .value_name("WHEN")
+            .possible_values(&["auto", "always", "never"])
+            .default_value("auto")
+            .global(true));
 
     if include_internal {
         app = app.subcommand(SubCommand::with_name("gen-completions")
@@ -140,6 +154,9 @@ pub async fn run() {
     let mut app = build_cli(true);
     let matches = app.clone().get_matches();
 
+    set_color_pref(matches.value_of("color").unwrap());
+    init_logging();
+
     handle_command!(apply, matches);
     handle_command!("apply-local", apply_local, matches);
     handle_command!(build, matches);
@@ -221,3 +238,29 @@ fn gen_help_markdown() {
         println!("</div></pre>");
     }
 }
+
+fn set_color_pref(cli: &str) {
+    if cli != "auto" {
+        clicolors_control::set_colors_enabled(cli == "always");
+    }
+}
+
+fn init_logging() {
+    if env::var("RUST_LOG").is_err() {
+        // HACK
+        env::set_var("RUST_LOG", "info")
+    }
+
+    // make env_logger conform to our detection logic
+    let style = if clicolors_control::colors_enabled() {
+        WriteStyle::Always
+    } else {
+        WriteStyle::Never
+    };
+
+    env_logger::builder()
+        .format_timestamp(None)
+        .format_module_path(false)
+        .write_style(style)
+        .init();
+}
diff --git a/src/main.rs b/src/main.rs
index 5d62828..ae9ba73 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,7 +1,5 @@
 #![deny(unused_must_use)]
 
-use std::env;
-
 mod nix;
 mod cli;
 mod command;
@@ -12,17 +10,5 @@ mod util;
 
 #[tokio::main]
 async fn main() {
-    init_logging();
     cli::run().await;
 }
-
-fn init_logging() {
-    if env::var("RUST_LOG").is_err() {
-        // HACK
-        env::set_var("RUST_LOG", "info")
-    }
-    env_logger::builder()
-        .format_timestamp(None)
-        .format_module_path(false)
-        .init();
-}