forked from cloudfoundry/cli
-
Notifications
You must be signed in to change notification settings - Fork 0
/
io.go
88 lines (73 loc) · 2.09 KB
/
io.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
package io
import (
"bytes"
"io"
"os"
"runtime"
"strings"
"github.com/fatih/color"
"github.com/mattn/go-colorable"
)
func SimulateStdin(input string, block func(r io.Reader)) {
reader, writer := io.Pipe()
go func() {
writer.Write([]byte(input))
defer writer.Close()
}()
block(reader)
}
func CaptureOutput(block func()) []string {
oldSTDOUT := os.Stdout
r, w, err := os.Pipe()
if err != nil {
panic(err)
}
os.Stdout = w
defer func() {
os.Stdout = oldSTDOUT
}()
//////
// We use fmt.Fprintf() to write to the "github.com/fatih/color".Output file
// to get colors on Windows machines.
// That variable gets initialized with a reference to os.Stdout when that library is imported.
// That means that when we muck with os.Stdout above, it doesn't get reflected in
// the printing code for windows.
// Instead, we can just redeclare that color.Output variable with a colorable version of our
// redirect pipe.
if runtime.GOOS == "windows" {
color.Output = colorable.NewColorable(w)
}
//////
doneWriting := make(chan bool)
result := make(chan []string)
go captureOutputAsynchronously(doneWriting, result, r)
block()
w.Close()
doneWriting <- true
return <-result
}
/*
The reason we're doing is that you can't write an infinite amount of bytes into a pipe.
On some platforms, the limit is fairly high; on other platforms, the limit is infuriatingly small
(looking at you, Windows). To counteract this, we need to read in a goroutine from one end of
the pipe and return the result across a channel.
*/
func captureOutputAsynchronously(doneWriting <-chan bool, result chan<- []string, reader io.Reader) {
var readingString string
for {
var buf bytes.Buffer
io.Copy(&buf, reader)
readingString += buf.String()
_, ok := <-doneWriting
if ok {
// there is no guarantee that the writer did not
// write more in between the read above and reading from this channel
// so we absolute must read once more if we want all the bytes
var buf bytes.Buffer
io.Copy(&buf, reader)
readingString += buf.String()
break
}
}
result <- strings.Split(readingString, "\n")
}