Fix command-line colorization on Windows with VT100-compatible environments

michaelpq · michaelpq · commit 3b77dce863d9 · 2020-03-02T15:45:34.000+09:00
When setting PG_COLOR to "always" or "auto" in a Windows terminal VT100-compatible, the colorization output was not showing up correctly because it is necessary to update the console's output handling mode. This fix allows to detect automatically if the environment is compatible with VT100. Hence, PG_COLOR=auto is able to detect and handle both compatible and non-compatible environments. The behavior of PG_COLOR=always remains unchanged, as it enforces the use of colorized output even if the environment does not allow it. This fix is based on an initial suggestion from Thomas Munro. Reported-by: Haiying Tang Author: Juan José Santamaría Flecha Reviewed-by: Michail Nikolaev, Michael Paquier, Haiying Tang Discussion: https://postgr.es/m/16108-134692e97146b7bc@postgresql.org Backpatch-through: 12
diff --git a/src/common/logging.c b/src/common/logging.c
@@ -32,6 +32,43 @@ static const char *sgr_locus = NULL;
 #define ANSI_ESCAPE_FMT "\x1b[%sm"
 #define ANSI_ESCAPE_RESET "\x1b[0m"
 
+#ifdef WIN32
+
+#ifndef ENABLE_VIRTUAL_TERMINAL_PROCESSING
+#define ENABLE_VIRTUAL_TERMINAL_PROCESSING 0x0004
+#endif
+
+/*
+ * Attempt to enable VT100 sequence processing for colorization on Windows.
+ * If current environment is not VT100-compatible or if this mode could not
+ * be enabled, return false.
+ */
+static bool
+enable_vt_processing(void)
+{
+	/* Check stderr */
+	HANDLE		hOut = GetStdHandle(STD_ERROR_HANDLE);
+	DWORD		dwMode = 0;
+
+	if (hOut == INVALID_HANDLE_VALUE)
+		return false;
+
+	/*
+	 * Look for the current console settings and check if VT100 is already
+	 * enabled.
+	 */
+	if (!GetConsoleMode(hOut, &dwMode))
+		return false;
+	if ((dwMode & ENABLE_VIRTUAL_TERMINAL_PROCESSING) != 0)
+		return true;
+
+	dwMode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING;
+	if (!SetConsoleMode(hOut, dwMode))
+		return false;
+	return true;
+}
+#endif							/* WIN32 */
+
 /*
  * This should be called before any output happens.
  */
@@ -40,6 +77,17 @@ pg_logging_init(const char *argv0)
 {
 	const char *pg_color_env = getenv("PG_COLOR");
 	bool		log_color = false;
+	bool		color_terminal = isatty(fileno(stderr));
+
+#ifdef WIN32
+
+	/*
+	 * On Windows, check if environment is VT100-compatible if using a
+	 * terminal.
+	 */
+	if (color_terminal)
+		color_terminal = enable_vt_processing();
+#endif
 
 	/* usually the default, but not on Windows */
 	setvbuf(stderr, NULL, _IONBF, 0);
@@ -50,7 +98,7 @@ pg_logging_init(const char *argv0)
 	if (pg_color_env)
 	{
 		if (strcmp(pg_color_env, "always") == 0 ||
-			(strcmp(pg_color_env, "auto") == 0 && isatty(fileno(stderr))))
+			(strcmp(pg_color_env, "auto") == 0 && color_terminal))
 			log_color = true;
 	}
 
