From ead727c274a5545af70bbca6b1b887c7055660fe Mon Sep 17 00:00:00 2001
From: Ruben Bridgewater <ruben@bridgewater.de>
Date: Tue, 16 Jan 2018 13:20:13 +0100
Subject: [PATCH] tty: add getColorDepth function

Right now it is very difficult to determine if a terminal supports
colors or not. This function adds this functionality by detecting
environment variables and checking process.

Backport-PR-URL: https://github.com/nodejs/node/pull/19230
PR-URL: https://github.com/nodejs/node/pull/17615
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
---
 doc/api/assert.md                         |  3 +
 doc/api/tty.md                            | 26 ++++++++
 lib/internal/errors.js                    | 21 +++++--
 lib/tty.js                                | 72 +++++++++++++++++++++++
 test/parallel/test-assert.js              |  9 ++-
 test/parallel/test-tty-get-color-depth.js | 52 ++++++++++++++++
 6 files changed, 175 insertions(+), 8 deletions(-)
 create mode 100644 test/parallel/test-tty-get-color-depth.js

diff --git a/doc/api/assert.md b/doc/api/assert.md
index 68b01350dfdee6..0698f54c23404f 100644
--- a/doc/api/assert.md
+++ b/doc/api/assert.md
@@ -65,6 +65,9 @@ AssertionError [ERR_ASSERTION]: Input A expected to deepStrictEqual input B:
   ]
 ```
 
+To deactivate the colors, use the `NODE_DISABLE_COLORS` environment variable.
+Please note that this will also deactivate the colors in the REPL.
+
 ## Legacy mode
 
 > Stability: 0 - Deprecated: Use strict mode instead.
diff --git a/doc/api/tty.md b/doc/api/tty.md
index 8b757c0f02751a..ce6dbae8fa6e7d 100644
--- a/doc/api/tty.md
+++ b/doc/api/tty.md
@@ -121,6 +121,32 @@ added: v0.7.7
 A `number` specifying the number of rows the TTY currently has. This property
 is updated whenever the `'resize'` event is emitted.
 
+### writeStream.getColorDepth([env])
+<!-- YAML
+added: REPLACEME
+-->
+
+* `env` {object} A object containing the environment variables to check.
+  Defaults to `process.env`.
+* Returns: {number}
+
+Returns:
+* 1 for 2,
+* 4 for 16,
+* 8 for 256,
+* 24 for 16,777,216
+colors supported.
+
+Use this to determine what colors the terminal supports. Due to the nature of
+colors in terminals it is possible to either have false positives or false
+negatives. It depends on process information and the environment variables that
+may lie about what terminal is used.
+To enforce a specific behavior without relying on `process.env` it is possible
+to pass in an object with different settings.
+
+Use the `NODE_DISABLE_COLORS` environment variable to enforce this function to
+always return 1.
+
 ## tty.isatty(fd)
 <!-- YAML
 added: v0.5.8
diff --git a/lib/internal/errors.js b/lib/internal/errors.js
index 02802746951fad..13f6dd38cee01d 100644
--- a/lib/internal/errors.js
+++ b/lib/internal/errors.js
@@ -13,6 +13,10 @@
 const kCode = Symbol('code');
 const messages = new Map();
 
+var green = '';
+var red = '';
+var white = '';
+
 const {
   UV_EAI_MEMORY,
   UV_EAI_NODATA,
@@ -90,7 +94,7 @@ function createErrDiff(actual, expected, operator) {
   const expectedLines = util
     .inspect(expected, { compact: false }).split('\n');
   const msg = `Input A expected to ${operator} input B:\n` +
-        '\u001b[32m+ expected\u001b[39m \u001b[31m- actual\u001b[39m';
+        `${green}+ expected${white} ${red}- actual${white}`;
   const skippedMsg = ' ... Lines skipped';
 
   // Remove all ending lines that match (this optimizes the output for
@@ -136,7 +140,7 @@ function createErrDiff(actual, expected, operator) {
         printedLines++;
       }
       lastPos = i;
-      other += `\n\u001b[32m+\u001b[39m ${expectedLines[i]}`;
+      other += `\n${green}+${white} ${expectedLines[i]}`;
       printedLines++;
     // Only extra actual lines exist
     } else if (expectedLines.length < i + 1) {
@@ -152,7 +156,7 @@ function createErrDiff(actual, expected, operator) {
         printedLines++;
       }
       lastPos = i;
-      res += `\n\u001b[31m-\u001b[39m ${actualLines[i]}`;
+      res += `\n${red}-${white} ${actualLines[i]}`;
       printedLines++;
     // Lines diverge
     } else if (actualLines[i] !== expectedLines[i]) {
@@ -168,8 +172,8 @@ function createErrDiff(actual, expected, operator) {
         printedLines++;
       }
       lastPos = i;
-      res += `\n\u001b[31m-\u001b[39m ${actualLines[i]}`;
-      other += `\n\u001b[32m+\u001b[39m ${expectedLines[i]}`;
+      res += `\n${red}-${white} ${actualLines[i]}`;
+      other += `\n${green}+${white} ${expectedLines[i]}`;
       printedLines += 2;
     // Lines are identical
     } else {
@@ -205,6 +209,13 @@ class AssertionError extends Error {
     if (message != null) {
       super(message);
     } else {
+      if (util_ === null &&
+          process.stdout.isTTY &&
+          process.stdout.getColorDepth() !== 1) {
+        green = '\u001b[32m';
+        white = '\u001b[39m';
+        red = '\u001b[31m';
+      }
       const util = lazyUtil();
       if (actual && actual.stack && actual instanceof Error)
         actual = `${actual.name}: ${actual.message}`;
diff --git a/lib/tty.js b/lib/tty.js
index 29440d6d96e3d5..8b6ce251c7c3ff 100644
--- a/lib/tty.js
+++ b/lib/tty.js
@@ -26,6 +26,14 @@ const net = require('net');
 const { TTY, isTTY } = process.binding('tty_wrap');
 const errors = require('internal/errors');
 const readline = require('readline');
+const { release } = require('os');
+
+const OSRelease = release().split('.');
+
+const COLORS_2 = 1;
+const COLORS_16 = 4;
+const COLORS_256 = 8;
+const COLORS_16m = 24;
 
 function isatty(fd) {
   return Number.isInteger(fd) && fd >= 0 && isTTY(fd);
@@ -90,6 +98,70 @@ inherits(WriteStream, net.Socket);
 
 WriteStream.prototype.isTTY = true;
 
+WriteStream.prototype.getColorDepth = function(env = process.env) {
+  if (env.NODE_DISABLE_COLORS || env.TERM === 'dumb' && !env.COLORTERM) {
+    return COLORS_2;
+  }
+
+  if (process.platform === 'win32') {
+    // Windows 10 build 10586 is the first Windows release that supports 256
+    // colors. Windows 10 build 14931 is the first release that supports
+    // 16m/TrueColor.
+    if (+OSRelease[0] >= 10) {
+      const build = +OSRelease[2];
+      if (build >= 14931)
+        return COLORS_16m;
+      if (build >= 10586)
+        return COLORS_256;
+    }
+
+    return COLORS_16;
+  }
+
+  if (env.TMUX) {
+    return COLORS_256;
+  }
+
+  if (env.CI) {
+    if ('TRAVIS' in env || 'CIRCLECI' in env || 'APPVEYOR' in env ||
+      'GITLAB_CI' in env || env.CI_NAME === 'codeship') {
+      return COLORS_256;
+    }
+    return COLORS_2;
+  }
+
+  if ('TEAMCITY_VERSION' in env) {
+    return /^(9\.(0*[1-9]\d*)\.|\d{2,}\.)/.test(env.TEAMCITY_VERSION) ?
+      COLORS_16 : COLORS_2;
+  }
+
+  switch (env.TERM_PROGRAM) {
+    case 'iTerm.app':
+      if (!env.TERM_PROGRAM_VERSION ||
+        /^[0-2]\./.test(env.TERM_PROGRAM_VERSION)) {
+        return COLORS_256;
+      }
+      return COLORS_16m;
+    case 'HyperTerm':
+    case 'Hyper':
+    case 'MacTerm':
+      return COLORS_16m;
+    case 'Apple_Terminal':
+      return COLORS_256;
+  }
+
+  if (env.TERM) {
+    if (/^xterm-256/.test(env.TERM))
+      return COLORS_256;
+    if (/^screen|^xterm|^vt100|color|ansi|cygwin|linux/i.test(env.TERM))
+      return COLORS_16;
+  }
+
+  if (env.COLORTERM)
+    return COLORS_16;
+
+  return COLORS_2;
+};
 
 WriteStream.prototype._refreshSize = function() {
   var oldCols = this.columns;
diff --git a/test/parallel/test-assert.js b/test/parallel/test-assert.js
index fa4ee21a6bd202..3f92e1035de1e3 100644
--- a/test/parallel/test-assert.js
+++ b/test/parallel/test-assert.js
@@ -763,10 +763,13 @@ common.expectsError(
   );
 
   // Test error diffs
+  const colors = process.stdout.isTTY && process.stdout.getColorDepth() > 1;
   const start = 'Input A expected to deepStrictEqual input B:';
-  const actExp = '\u001b[32m+ expected\u001b[39m \u001b[31m- actual\u001b[39m';
-  const plus = '\u001b[32m+\u001b[39m';
-  const minus = '\u001b[31m-\u001b[39m';
+  const actExp = colors ?
+    '\u001b[32m+ expected\u001b[39m \u001b[31m- actual\u001b[39m' :
+    '+ expected - actual';
+  const plus = colors ? '\u001b[32m+\u001b[39m' : '+';
+  const minus = colors ? '\u001b[31m-\u001b[39m' : '-';
   let message = [
     start,
     `${actExp} ... Lines skipped`,
diff --git a/test/parallel/test-tty-get-color-depth.js b/test/parallel/test-tty-get-color-depth.js
new file mode 100644
index 00000000000000..a5a998b6a3ab78
--- /dev/null
+++ b/test/parallel/test-tty-get-color-depth.js
@@ -0,0 +1,52 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert').strict;
+/* eslint-disable no-restricted-properties */
+const { openSync } = require('fs');
+const tty = require('tty');
+
+const { WriteStream } = require('tty');
+
+// Do our best to grab a tty fd.
+function getTTYfd() {
+  const ttyFd = [0, 1, 2, 4, 5].find(tty.isatty);
+  if (ttyFd === undefined) {
+    try {
+      return openSync('/dev/tty');
+    } catch (e) {
+      // There aren't any tty fd's available to use.
+      return -1;
+    }
+  }
+  return ttyFd;
+}
+
+const fd = getTTYfd();
+
+// Give up if we did not find a tty
+if (fd === -1)
+  common.skip();
+
+const writeStream = new WriteStream(fd);
+
+let depth = writeStream.getColorDepth();
+
+assert.equal(typeof depth, 'number');
+assert(depth >= 1 && depth <= 24);
+
+// If the terminal does not support colors, skip the rest
+if (depth === 1)
+  common.skip();
+
+assert.notEqual(writeStream.getColorDepth({ TERM: 'dumb' }), depth);
+
+// Deactivate colors
+const tmp = process.env.NODE_DISABLE_COLORS;
+process.env.NODE_DISABLE_COLORS = 1;
+
+depth = writeStream.getColorDepth();
+
+assert.equal(depth, 1);
+
+process.env.NODE_DISABLE_COLORS = tmp;