Merge pull request cli#4146 from cli/force-tty

mislav · web-flow · commit af2aecd40bf6 · 2021-08-23T16:46:14.000+02:00
Add ability to force terminal-style output even when redirected
diff --git a/cmd/gh/main.go b/cmd/gh/main.go
@@ -61,6 +61,10 @@ func mainRun() exitCode {
 
 	cmdFactory := factory.New(buildVersion)
 	stderr := cmdFactory.IOStreams.ErrOut
+
+	if spec := os.Getenv("GH_FORCE_TTY"); spec != "" {
+		cmdFactory.IOStreams.ForceTerminal(spec)
+	}
 	if !cmdFactory.IOStreams.ColorEnabled() {
 		surveyCore.DisableColor = true
 	} else {
diff --git a/pkg/cmd/root/help_topic.go b/pkg/cmd/root/help_topic.go
@@ -62,6 +62,11 @@ var HelpTopics = map[string]map[string]string{
 			CLICOLOR_FORCE: set to a value other than "0" to keep ANSI colors in output
 			even when the output is piped.
 
+			GH_FORCE_TTY: set to any value to force terminal-style output even when the output is
+			redirected. When the value is a number, it is interpreted as the number of columns
+			available in the viewport. When the value is a percentage, it will be applied against
+			the number of columns available in the current viewport.
+
 			GH_NO_UPDATE_NOTIFIER: set to any value to disable update notifications. By default, gh
 			checks for new releases once every 24 hours and displays an upgrade notice on standard
 			error if a newer version was found.
diff --git a/pkg/iostreams/iostreams.go b/pkg/iostreams/iostreams.go
@@ -2,6 +2,7 @@ package iostreams
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -40,6 +41,8 @@ type IOStreams struct {
 	stdoutIsTTY       bool
 	stderrTTYOverride bool
 	stderrIsTTY       bool
+	termWidthOverride int
+	ttySize           func() (int, int, error)
 
 	pagerCommand string
 	pagerProcess *os.Process
@@ -232,6 +235,10 @@ func (s *IOStreams) StopProgressIndicator() {
 }
 
 func (s *IOStreams) TerminalWidth() int {
+	if s.termWidthOverride > 0 {
+		return s.termWidthOverride
+	}
+
 	defaultWidth := 80
 	out := s.Out
 	if s.originalOut != nil {
@@ -259,6 +266,28 @@ func (s *IOStreams) TerminalWidth() int {
 	return defaultWidth
 }
 
+func (s *IOStreams) ForceTerminal(spec string) {
+	s.colorEnabled = !EnvColorDisabled()
+	s.SetStdoutTTY(true)
+
+	if w, err := strconv.Atoi(spec); err == nil {
+		s.termWidthOverride = w
+		return
+	}
+
+	ttyWidth, _, err := s.ttySize()
+	if err != nil {
+		return
+	}
+	s.termWidthOverride = ttyWidth
+
+	if strings.HasSuffix(spec, "%") {
+		if p, err := strconv.Atoi(spec[:len(spec)-1]); err == nil {
+			s.termWidthOverride = int(float64(s.termWidthOverride) * (float64(p) / 100))
+		}
+	}
+}
+
 func (s *IOStreams) ColorScheme() *ColorScheme {
 	return NewColorScheme(s.ColorEnabled(), s.ColorSupport256())
 }
@@ -297,6 +326,7 @@ func System() *IOStreams {
 		colorEnabled: EnvColorForced() || (!EnvColorDisabled() && stdoutIsTTY),
 		is256enabled: Is256ColorSupported(),
 		pagerCommand: os.Getenv("PAGER"),
+		ttySize:      ttySize,
 	}
 
 	if stdoutIsTTY && stderrIsTTY {
@@ -317,6 +347,9 @@ func Test() (*IOStreams, *bytes.Buffer, *bytes.Buffer, *bytes.Buffer) {
 		In:     ioutil.NopCloser(in),
 		Out:    out,
 		ErrOut: errOut,
+		ttySize: func() (int, int, error) {
+			return -1, -1, errors.New("ttySize not implemented in tests")
+		},
 	}, in, out, errOut
 }
 
@@ -331,6 +364,7 @@ func isCygwinTerminal(w io.Writer) bool {
 	return false
 }
 
+// terminalSize measures the viewport of the terminal that the output stream is connected to
 func terminalSize(w io.Writer) (int, int, error) {
 	if f, isFile := w.(*os.File); isFile {
 		return term.GetSize(int(f.Fd()))
diff --git a/pkg/iostreams/iostreams_test.go b/pkg/iostreams/iostreams_test.go
@@ -0,0 +1,68 @@
+package iostreams
+
+import (
+	"errors"
+	"testing"
+)
+
+func TestIOStreams_ForceTerminal(t *testing.T) {
+	tests := []struct {
+		name      string
+		iostreams *IOStreams
+		arg       string
+		wantTTY   bool
+		wantWidth int
+	}{
+		{
+			name:      "explicit width",
+			iostreams: &IOStreams{},
+			arg:       "72",
+			wantTTY:   true,
+			wantWidth: 72,
+		},
+		{
+			name: "measure width",
+			iostreams: &IOStreams{
+				ttySize: func() (int, int, error) {
+					return 72, 0, nil
+				},
+			},
+			arg:       "true",
+			wantTTY:   true,
+			wantWidth: 72,
+		},
+		{
+			name: "measure width fails",
+			iostreams: &IOStreams{
+				ttySize: func() (int, int, error) {
+					return -1, -1, errors.New("ttySize sabotage!")
+				},
+			},
+			arg:       "true",
+			wantTTY:   true,
+			wantWidth: 80,
+		},
+		{
+			name: "apply percentage",
+			iostreams: &IOStreams{
+				ttySize: func() (int, int, error) {
+					return 72, 0, nil
+				},
+			},
+			arg:       "50%",
+			wantTTY:   true,
+			wantWidth: 36,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.iostreams.ForceTerminal(tt.arg)
+			if isTTY := tt.iostreams.IsStdoutTTY(); isTTY != tt.wantTTY {
+				t.Errorf("IOStreams.IsStdoutTTY() = %v, want %v", isTTY, tt.wantTTY)
+			}
+			if tw := tt.iostreams.TerminalWidth(); tw != tt.wantWidth {
+				t.Errorf("IOStreams.TerminalWidth() = %v, want %v", tw, tt.wantWidth)
+			}
+		})
+	}
+}
diff --git a/pkg/iostreams/tty_size.go b/pkg/iostreams/tty_size.go
@@ -0,0 +1,19 @@
+//+build !windows
+
+package iostreams
+
+import (
+	"os"
+
+	"golang.org/x/term"
+)
+
+// ttySize measures the size of the controlling terminal for the current process
+func ttySize() (int, int, error) {
+	f, err := os.Open("/dev/tty")
+	if err != nil {
+		return -1, -1, err
+	}
+	defer f.Close()
+	return term.GetSize(int(f.Fd()))
+}
diff --git a/pkg/iostreams/tty_size_windows.go b/pkg/iostreams/tty_size_windows.go
@@ -0,0 +1,16 @@
+package iostreams
+
+import (
+	"os"
+
+	"golang.org/x/term"
+)
+
+func ttySize() (int, int, error) {
+	f, err := os.Open("CONOUT$")
+	if err != nil {
+		return -1, -1, err
+	}
+	defer f.Close()
+	return term.GetSize(int(f.Fd()))
+}
