From ecdbf8b46e794aec4245ffdeca9fda3b43b7e6b8 Mon Sep 17 00:00:00 2001
From: Lorenzo Gabriele <lorenzolespaul@gmail.com>
Date: Sat, 6 Jan 2024 22:37:06 +0100
Subject: [PATCH] Enable colored output for Scala 2 (#2950)

Fixes #2947

Pull Request: https://github.com/com-lihaoyi/mill/pull/2950
---
 .../core/src/Main.scala                       | 12 +++++++
 .../src/mill/scalalib/HelloWorldTests.scala   | 33 +++++++++++++++++--
 .../mill/scalalib/worker/ZincWorkerImpl.scala |  7 ++++
 3 files changed, 49 insertions(+), 3 deletions(-)
 create mode 100644 scalalib/test/resources/hello-world-color-output/core/src/Main.scala

diff --git a/scalalib/test/resources/hello-world-color-output/core/src/Main.scala b/scalalib/test/resources/hello-world-color-output/core/src/Main.scala
new file mode 100644
index 00000000000..3e878e0f3f1
--- /dev/null
+++ b/scalalib/test/resources/hello-world-color-output/core/src/Main.scala
@@ -0,0 +1,12 @@
+package example
+
+trait Show[A]
+
+object Show {
+  implicit def option[A](implicit s: Show[A]): Show[Option[A]] = ???
+}
+
+object Example {
+  def main(args: Array[String]): Unit =
+    println(implicitly[Show[Option[String]]])
+}
diff --git a/scalalib/test/src/mill/scalalib/HelloWorldTests.scala b/scalalib/test/src/mill/scalalib/HelloWorldTests.scala
index cdc3e0d71fb..873536a6bd5 100644
--- a/scalalib/test/src/mill/scalalib/HelloWorldTests.scala
+++ b/scalalib/test/src/mill/scalalib/HelloWorldTests.scala
@@ -1,6 +1,6 @@
 package mill.scalalib
 
-import java.io.ByteArrayOutputStream
+import java.io.{ByteArrayOutputStream, PrintStream}
 import java.util.jar.JarFile
 import scala.jdk.CollectionConverters._
 import scala.util.{Properties, Using}
@@ -283,6 +283,16 @@ object HelloWorldTests extends TestSuite {
     }
   }
 
+  object HelloWorldColorOutput extends HelloBase {
+    object core extends ScalaModule {
+      def scalaVersion = scala213Version
+
+      override def scalacOptions = super.scalacOptions() ++ Seq(
+        "-Vimplicits"
+      )
+    }
+  }
+
   object HelloScalacheck extends HelloBase {
     object foo extends ScalaModule {
       def scalaVersion = scala212Version
@@ -420,9 +430,10 @@ object HelloWorldTests extends TestSuite {
       m: TestUtil.BaseModule,
       resourcePath: os.Path = resourcePath,
       env: Map[String, String] = Evaluator.defaultEnv,
-      debug: Boolean = false
+      debug: Boolean = false,
+      errStream: PrintStream = System.err
   )(t: TestEvaluator => T)(implicit tp: TestPath): T = {
-    val eval = new TestEvaluator(m, env = env, debugEnabled = debug)
+    val eval = new TestEvaluator(m, env = env, debugEnabled = debug, errStream = errStream)
     os.remove.all(m.millSourcePath)
     os.remove.all(eval.outPath)
     os.makeDir.all(m.millSourcePath / os.up)
@@ -1189,6 +1200,22 @@ object HelloWorldTests extends TestSuite {
         assert(evalCount > 0)
       }
     }
+    "color-output" - {
+      val errStream = new ByteArrayOutputStream()
+      workspaceTest(
+        HelloWorldColorOutput,
+        os.pwd / "scalalib" / "test" / "resources" / "hello-world-color-output",
+        errStream = new PrintStream(errStream, true)
+      ) { eval =>
+        val Left(Result.Failure("Compilation failed", _)) =
+          eval.apply(HelloWorldColorOutput.core.compile)
+        val output = errStream.toString
+        assert(output.contains(s"${Console.RED}!${Console.RESET}${Console.BLUE}I"))
+        assert(output.contains(
+          s"${Console.GREEN}example.Show[scala.Option[java.lang.String]]${Console.RESET}"
+        ))
+      }
+    }
 
     "scalacheck" - workspaceTest(
       HelloScalacheck,
diff --git a/scalalib/worker/src/mill/scalalib/worker/ZincWorkerImpl.scala b/scalalib/worker/src/mill/scalalib/worker/ZincWorkerImpl.scala
index 3a93433e86a..1e18e823e1c 100644
--- a/scalalib/worker/src/mill/scalalib/worker/ZincWorkerImpl.scala
+++ b/scalalib/worker/src/mill/scalalib/worker/ZincWorkerImpl.scala
@@ -534,7 +534,10 @@ class ZincWorkerImpl(
       stampReader = Stamps.timeWrapBinaryStamps(converter)
     )
 
+    val scalaColorProp = "scala.color"
+    val previousScalaColor = sys.props(scalaColorProp)
     try {
+      sys.props(scalaColorProp) = if (ctx.log.colored) "true" else "false"
       val newResult = ic.compile(
         in = inputs,
         logger = logger
@@ -557,6 +560,10 @@ class ZincWorkerImpl(
     } finally {
       reporter.foreach(r => sources.foreach(r.fileVisited(_)))
       reporter.foreach(_.finish())
+      previousScalaColor match {
+        case null => sys.props.remove(scalaColorProp)
+        case v => sys.props(scalaColorProp) = previousScalaColor
+      }
     }
   }