From d696bd113bc64cd879c83e6d1def94b2fcb3b6c7 Mon Sep 17 00:00:00 2001
From: Jeremie Dimino <jdimino@janestreet.com>
Date: Thu, 23 Feb 2017 12:12:02 +0000
Subject: [PATCH] Hack to force OCaml to use colors when possible

---
 src/ansi_color.ml  | 24 ++++++++++++++++++++++++
 src/ansi_color.mli |  2 ++
 src/context.ml     |  4 +++-
 src/future.ml      | 26 +++++++++++++++-----------
 4 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/src/ansi_color.ml b/src/ansi_color.ml
index 249c29cb..9b6c9cb7 100644
--- a/src/ansi_color.ml
+++ b/src/ansi_color.ml
@@ -107,3 +107,27 @@ let strip str =
       | _   -> skip (i + 1)
   in
   loop 0
+
+let stderr_supports_colors = lazy(
+  not Sys.win32        &&
+  Unix.(isatty stderr) &&
+  match Sys.getenv "TERM" with
+  | exception Not_found -> false
+  | "dumb" -> false
+  | _ -> true
+)
+
+(* We redirect the output of all commands, so by default the compiler will disable
+   colors. Since we support colors in the output of commands, we force it via OCAMLPARAM
+   if stderr supports colors. *)
+let setup_env_for_ocaml_colors = lazy(
+  if Lazy.force stderr_supports_colors then begin
+    let value =
+      match Sys.getenv "OCAMLPARAM" with
+      | exception Not_found -> "color=always,_"
+      | s -> "color=always," ^ s
+    in
+    Unix.putenv "OCAMLPARAM" value
+  end
+)
+
diff --git a/src/ansi_color.mli b/src/ansi_color.mli
index 202e7430..e895ecc3 100644
--- a/src/ansi_color.mli
+++ b/src/ansi_color.mli
@@ -1,2 +1,4 @@
 val colorize : key:string -> string -> string
 val strip : string -> string
+val stderr_supports_colors : bool Lazy.t
+val setup_env_for_ocaml_colors : unit Lazy.t
diff --git a/src/context.ml b/src/context.ml
index 14cf2d95..aafce7af 100644
--- a/src/context.ml
+++ b/src/context.ml
@@ -210,7 +210,9 @@ let create ~(kind : Kind.t) ~path ~env =
   all_known := String_map.add !all_known ~key:name ~data:t;
   return t
 
-let initial_env = lazy (Unix.environment ())
+let initial_env = lazy (
+  Lazy.force Ansi_color.setup_env_for_ocaml_colors;
+  Unix.environment ())
 
 let default = lazy (
   let env = Lazy.force initial_env in
diff --git a/src/future.ml b/src/future.ml
index 1d803f1c..8bfefb30 100644
--- a/src/future.ml
+++ b/src/future.ml
@@ -214,17 +214,20 @@ module Scheduler = struct
     | None -> s
     | Some dir -> sprintf "(cd %s && %s)" dir s
 
-  let strip_colors_for_stderr =
-    let strip = lazy (
-      Sys.win32                  ||
-      not (Unix.(isatty stderr)) ||
-      match Sys.getenv "TERM" with
-      | exception Not_found -> true
-      | "dumb" -> true
-      | _ -> false
-    ) in
-    fun s ->
-      if Lazy.force strip then Ansi_color.strip s else s
+  let stderr_supports_colors = lazy(
+    not Sys.win32        &&
+    Unix.(isatty stderr) &&
+    match Sys.getenv "TERM" with
+    | exception Not_found -> false
+    | "dumb" -> false
+    | _ -> true
+  )
+
+  let strip_colors_for_stderr s =
+    if Lazy.force Ansi_color.stderr_supports_colors then
+      s
+    else
+      Ansi_color.strip s
 
   type running_job =
     { id              : int
@@ -375,6 +378,7 @@ module Scheduler = struct
       go_rec cwd log t
 
   let go ?log t =
+    Lazy.force Ansi_color.setup_env_for_ocaml_colors;
     let cwd = Sys.getcwd () in
     go_rec cwd log t
 end