From 70ac6ec186bc249f2dfa1e9511c5a6d796509eec Mon Sep 17 00:00:00 2001
From: George Dietrich <dev@dietrich.pub>
Date: Sun, 15 Dec 2024 13:41:24 -0500
Subject: [PATCH] More robust `console` color support logic
 (athena-framework/athena#488)

* Improved spec coverage
---
 .../abstract_question_helper_test_case.cr     |  2 +-
 src/components/console/spec/output/io_spec.cr | 79 +++++++++++++++++++
 src/components/console/spec/spec_helper.cr    | 14 ++++
 src/components/console/src/output/io.cr       | 20 ++++-
 4 files changed, 111 insertions(+), 4 deletions(-)

diff --git a/src/components/console/spec/helper/abstract_question_helper_test_case.cr b/src/components/console/spec/helper/abstract_question_helper_test_case.cr
index ba9852a61..863f63ca8 100644
--- a/src/components/console/spec/helper/abstract_question_helper_test_case.cr
+++ b/src/components/console/spec/helper/abstract_question_helper_test_case.cr
@@ -4,7 +4,7 @@ abstract struct AbstractQuestionHelperTest < ASPEC::TestCase
   def initialize
     @helper_set = ACON::Helper::HelperSet.new ACON::Helper::Formatter.new
 
-    @output = ACON::Output::IO.new IO::Memory.new
+    @output = ACON::Output::IO.new IO::Memory.new, decorated: false
   end
 
   protected def with_input(data : String, interactive : Bool = true, & : ACON::Input::Interface -> Nil) : Nil
diff --git a/src/components/console/spec/output/io_spec.cr b/src/components/console/spec/output/io_spec.cr
index cfebd57f9..0700ba541 100644
--- a/src/components/console/spec/output/io_spec.cr
+++ b/src/components/console/spec/output/io_spec.cr
@@ -16,4 +16,83 @@ struct IOTest < ASPEC::TestCase
     output.puts "foo"
     output.to_s.should eq "foo#{EOL}"
   end
+
+  def test_decorated_dumb_term : Nil
+    with_isolated_env do
+      ENV["TERM"] = "dumb"
+      ACON::Output::IO.new(@io).decorated?.should be_false
+    end
+  end
+
+  def test_decorated_no_color : Nil
+    with_isolated_env do
+      ENV["NO_COLOR"] = "true"
+      ENV["COLORTERM"] = "truecolor"
+      ACON::Output::IO.new(@io).decorated?.should be_false
+    end
+  end
+
+  def test_decorated_no_color_empty : Nil
+    with_isolated_env do
+      ENV["NO_COLOR"] = ""
+      ENV["COLORTERM"] = "truecolor"
+      ACON::Output::IO.new(@io).decorated?.should be_true
+    end
+  end
+
+  def test_decorated_force_color : Nil
+    with_isolated_env do
+      ENV["FORCE_COLOR"] = "true"
+      ACON::Output::IO.new(@io).decorated?.should be_true
+    end
+  end
+
+  def test_decorated_force_color_empty : Nil
+    with_isolated_env do
+      ENV["FORCE_COLOR"] = ""
+      ACON::Output::IO.new(@io).decorated?.should be_false
+    end
+  end
+
+  def test_decorated_supported_term : Nil
+    with_isolated_env do
+      ENV["TERM"] = "xterm-256color"
+      ACON::Output::IO.new(@io).decorated?.should be_true
+    end
+  end
+
+  def test_decorated_colorterm : Nil
+    with_isolated_env do
+      ENV["COLORTERM"] = "truecolor"
+      ACON::Output::IO.new(@io).decorated?.should be_true
+    end
+  end
+
+  def test_decorated_ansicon : Nil
+    with_isolated_env do
+      ENV["ANSICON"] = "1"
+      ACON::Output::IO.new(@io).decorated?.should be_true
+    end
+  end
+
+  def test_decorated_conemuansi : Nil
+    with_isolated_env do
+      ENV["ConEmuANSI"] = "ON"
+      ACON::Output::IO.new(@io).decorated?.should be_true
+    end
+  end
+
+  def test_decorated_term_program_hyper : Nil
+    with_isolated_env do
+      ENV["TERM_PROGRAM"] = "Hyper"
+      ACON::Output::IO.new(@io).decorated?.should be_true
+    end
+  end
+
+  def test_decorated_term_program_non_hyper : Nil
+    with_isolated_env do
+      ENV["TERM_PROGRAM"] = "WezTerm"
+      ACON::Output::IO.new(@io).decorated?.should be_false
+    end
+  end
 end
diff --git a/src/components/console/spec/spec_helper.cr b/src/components/console/spec/spec_helper.cr
index b427902ef..8060d7d1d 100644
--- a/src/components/console/spec/spec_helper.cr
+++ b/src/components/console/spec/spec_helper.cr
@@ -44,4 +44,18 @@ struct MockCommandLoader
   end
 end
 
+def with_isolated_env(&) : Nil
+  old_values = ENV.dup
+  begin
+    ENV.clear
+
+    yield
+  ensure
+    ENV.clear
+    old_values.each do |key, old_value|
+      ENV[key] = old_value
+    end
+  end
+end
+
 ASPEC.run_all
diff --git a/src/components/console/src/output/io.cr b/src/components/console/src/output/io.cr
index 6727e0b94..4ffcb3628 100644
--- a/src/components/console/src/output/io.cr
+++ b/src/components/console/src/output/io.cr
@@ -29,9 +29,23 @@ class Athena::Console::Output::IO < Athena::Console::Output
 
   private def has_color_support? : Bool
     # Respect https://no-color.org.
-    return false if "false" == ENV["NO_COLOR"]?
-    return true if "Hyper" == ENV["TERM_PROGRAM"]?
+    return false if ENV["NO_COLOR"]?.presence
 
-    @io.tty?
+    # Respect https://force-color.org.
+    return true if ENV["FORCE_COLOR"]?.presence
+
+    if "Hyper" == ENV["TERM_PROGRAM"]? ||
+       ENV.has_key?("COLORTERM") ||
+       ENV.has_key?("ANSICON") ||
+       "ON" == ENV["ConEmuANSI"]?
+      return true
+    end
+
+    return @io.tty? unless term = ENV["TERM"]?
+
+    return false if "dumb" == term
+
+    # See https://github.com/chalk/supports-color/blob/d4f413efaf8da045c5ab440ed418ef02dbb28bf1/index.js#L157
+    term.matches? /^((screen|xterm|vt100|vt220|putty|rxvt|ansi|cygwin|linux).*)|(.*-256(color)?(-bce)?)$/
   end
 end