From ef7fc6b790e9c6427083d93b4775563b83c47f9a Mon Sep 17 00:00:00 2001
From: Tom Payne <twpayne@gmail.com>
Date: Sat, 17 Sep 2022 00:49:03 +0200
Subject: [PATCH] Add EnableVirtualTerminalProcessing function

The existing EnableWindowsANSIConsole/RestoreWindowsConsole functions
have a number of limitations:
* They are only defined when termenv is built on Windows, and so require
  the user to use build tags, i.e. multiple source files, to control
  whether or not they are called and are not shown on
  https://pkg.go.dev/github.com/muesli/termenv by default.
* They are hardcoded to set the console mode of stdout, and so fail if
  stdout is not a terminal, e.g. when redirecting the output to a file
  or when run in a Go test.

This commit adds a EnableVirtualTerminalProcessing function with a
different API (to avoid breaking backwards compatibility) that is safe
to call on all platforms, takes an io.Writer as an argument (for output
flexibility). See the comments in the function for more details.
---
 termenv_test.go    | 13 +++++++++++++
 termenv_unix.go    |  9 +++++++++
 termenv_windows.go | 47 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 69 insertions(+)

diff --git a/termenv_test.go b/termenv_test.go
index 8b59d98..87f22ad 100644
--- a/termenv_test.go
+++ b/termenv_test.go
@@ -392,3 +392,16 @@ func TestCache(t *testing.T) {
 		t.Errorf("Expected cache to be active, got %t", o.cache)
 	}
 }
+
+func TestEnableVirtualTerminalProcessing(t *testing.T) {
+	// EnableVirtualTerminalProcessing should always return a non-nil
+	// restoreFunc, and in tests it should never return an error.
+	restoreFunc, err := EnableVirtualTerminalProcessing(NewOutput(os.Stdout))
+	if restoreFunc == nil || err != nil {
+		t.Fatalf("expected non-<nil>, <nil>, got %p, %v", restoreFunc, err)
+	}
+	// In tests, restoreFunc should never return an error.
+	if err := restoreFunc(); err != nil {
+		t.Fatalf("expected <nil>, got %v", err)
+	}
+}
diff --git a/termenv_unix.go b/termenv_unix.go
index 333a08b..778c354 100644
--- a/termenv_unix.go
+++ b/termenv_unix.go
@@ -5,6 +5,7 @@ package termenv
 
 import (
 	"fmt"
+	"io"
 	"strconv"
 	"strings"
 	"time"
@@ -273,3 +274,11 @@ func (o Output) termStatusReport(sequence int) (string, error) {
 	// fmt.Println("Rcvd", res[1:])
 	return res, nil
 }
+
+// EnableVirtualTerminalProcessing enables virtual terminal processing on
+// Windows for w and returns a function that restores w to its previous state.
+// On non-Windows platforms, or if w does not refer to a terminal, then it
+// returns a non-nil no-op function and no error.
+func EnableVirtualTerminalProcessing(w io.Writer) (func() error, error) {
+	return func() error { return nil }, nil
+}
diff --git a/termenv_windows.go b/termenv_windows.go
index 84e5c2e..1d9c618 100644
--- a/termenv_windows.go
+++ b/termenv_windows.go
@@ -4,6 +4,7 @@
 package termenv
 
 import (
+	"fmt"
 	"strconv"
 
 	"golang.org/x/sys/windows"
@@ -90,3 +91,49 @@ func RestoreWindowsConsole(mode uint32) error {
 
 	return windows.SetConsoleMode(handle, mode)
 }
+
+// EnableVirtualTerminalProcessing enables virtual terminal processing on
+// Windows for o and returns a function that restores o to its previous state.
+// On non-Windows platforms, or if o does not refer to a terminal, then it
+// returns a non-nil no-op function and no error.
+func EnableVirtualTerminalProcessing(o *Output) (restoreFunc func() error, err error) {
+	// There is nothing to restore until we set the console mode.
+	restoreFunc = func() error {
+		return nil
+	}
+
+	// If o is not a tty, then there is nothing to do.
+	tty := o.TTY()
+	if tty == nil {
+		return
+	}
+
+	// Get the current console mode. If there is an error, assume that o is not
+	// a terminal, discard the error, and return.
+	var mode uint32
+	if err2 := windows.GetConsoleMode(windows.Handle(tty.Fd()), &mode); err2 != nil {
+		return
+	}
+
+	// If virtual terminal processing is already set, then there is nothing to
+	// do and nothing to restore.
+	if mode&windows.ENABLE_VIRTUAL_TERMINAL_PROCESSING == windows.ENABLE_VIRTUAL_TERMINAL_PROCESSING {
+		return
+	}
+
+	// Enable virtual terminal processing. See
+	// https://docs.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences
+	if err2 := windows.SetConsoleMode(windows.Handle(tty.Fd()), mode|windows.ENABLE_VIRTUAL_TERMINAL_PROCESSING); err2 != nil {
+		err = fmt.Errorf("windows.SetConsoleMode: %w", err2)
+		return
+	}
+
+	// Set the restore function. We maintain a reference to the tty in the
+	// closure (rather than just its handle) to ensure that the tty is not
+	// closed by a finalizer.
+	restoreFunc = func() error {
+		return windows.SetConsoleMode(windows.Handle(tty.Fd()), mode)
+	}
+
+	return
+}