feat: add --color arg for cli (#34)

Signed-off-by: Alex Chi <[email protected]>
risinglightdb · Mar 30, 2022 · ecf60e9 · ecf60e9
1 parent b9a0cbe
commit ecf60e9
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 8 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - New test UI for sqllogictest binary
+- Support set console color on the test UI with `--color` parameter
 
 ## [0.3.0] - 2022-03-21
 

diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "sqllogictest"
-version = "0.3.2"
+version = "0.3.3"
 edition = "2021"
 description = "Sqllogictest parser and runner."
 license = "MIT OR Apache-2.0"
@@ -16,7 +16,7 @@ bin = ["anyhow", "console", "tokio-postgres", "env_logger", "glob", "clap", "tok
 [dependencies]
 anyhow = { version = "1", optional = true }
 async-trait = "0.1"
-clap = { version = "3", features = ["derive"], optional = true }
+clap = { version = "3", features = ["derive", "env"], optional = true }
 console = { version = "0.15", optional = true }
 env_logger = { version = "0.9", optional = true }
 futures-lite = "1"

diff --git a/src/main.rs b/src/main.rs
@@ -5,10 +5,24 @@ use std::time::Instant;
 
 use anyhow::{anyhow, Context, Result};
 use async_trait::async_trait;
-use clap::Parser;
+use clap::{ArgEnum, Parser};
 use console::style;
 use sqllogictest::{Control, Record};
 
+#[derive(Copy, Clone, Debug, PartialEq, ArgEnum)]
+#[must_use]
+pub enum Color {
+    Auto,
+    Always,
+    Never,
+}
+
+impl Default for Color {
+    fn default() -> Self {
+        Color::Auto
+    }
+}
+
 #[derive(Parser, Debug)]
 #[clap(about, version, author)]
 struct Opt {
@@ -40,6 +54,15 @@ struct Opt {
     /// The database password.
     #[clap(short = 'w', long, default_value = "postgres")]
     pass: String,
+
+    #[clap(
+        long,
+        arg_enum,
+        default_value_t,
+        value_name = "WHEN",
+        env = "CARGO_TERM_COLOR"
+    )]
+    color: Color,
 }
 
 #[tokio::main]
@@ -48,6 +71,18 @@ async fn main() -> Result<()> {
 
     let opt = Opt::parse();
 
+    match opt.color {
+        Color::Always => {
+            console::set_colors_enabled(true);
+            console::set_colors_enabled_stderr(true);
+        }
+        Color::Never => {
+            console::set_colors_enabled(false);
+            console::set_colors_enabled_stderr(false);
+        }
+        Color::Auto => {}
+    }
+
     let files = glob::glob(&opt.files).expect("failed to read glob pattern");
 
     let (client, connection) = tokio_postgres::Config::new()
@@ -107,7 +142,7 @@ async fn run_test_file(engine: Postgres, filename: impl AsRef<Path>) -> Result<(
     let mut begin_times = vec![];
     let mut did_pop = false;
 
-    print!("{} .. ", filename.to_string_lossy());
+    print!("{: <60} .. ", filename.to_string_lossy());
     flush_stdout().await?;
 
     begin_times.push(Instant::now());
@@ -118,10 +153,10 @@ async fn run_test_file(engine: Postgres, filename: impl AsRef<Path>) -> Result<(
         if *did_pop {
             // start a new line if the result is not immediately after the item
             print!(
-                "\n{}{} {} {} in {} ms",
-                "  ".repeat(time_stack.len()),
-                file,
+                "\n{}{} {: <54} .. {} in {} ms",
+                "| ".repeat(time_stack.len()),
                 style("[END]").blue().bold(),
+                file,
                 style("[OK]").green().bold(),
                 begin_time.elapsed().as_millis()
             );
@@ -147,7 +182,7 @@ async fn run_test_file(engine: Postgres, filename: impl AsRef<Path>) -> Result<(
                     println!();
                 }
                 did_pop = false;
-                print!("{}{} .. ", "  ".repeat(begin_times.len() - 1), file);
+                print!("{}{: <60} .. ", "| ".repeat(begin_times.len() - 1), file);
                 flush_stdout().await?;
             }
             Record::Control(Control::EndInclude(file)) => {