-
Notifications
You must be signed in to change notification settings - Fork 1
/
colorout.go
143 lines (126 loc) · 3.32 KB
/
colorout.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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
// Copyright 2016-2017 Joonas Kuorilehto.
// Command colorout colors and multiplexes output from tasks.
// Each task (external command) is executed concurrently using
// command shell and output is colored with a unique color per task.
package main
import (
"bytes"
"context"
"flag"
"fmt"
"io"
"log"
"os"
"os/exec"
"sync"
"github.com/fatih/color"
)
var colors = []*color.Color{
color.New(color.FgHiRed),
color.New(color.FgHiGreen),
color.New(color.FgHiYellow),
color.New(color.FgHiBlue),
color.New(color.FgHiMagenta),
color.New(color.FgHiCyan),
color.New(color.FgHiWhite),
color.New(color.FgRed, color.ReverseVideo),
color.New(color.FgGreen, color.ReverseVideo),
color.New(color.FgYellow, color.ReverseVideo),
color.New(color.FgBlue, color.ReverseVideo),
color.New(color.FgMagenta, color.ReverseVideo),
color.New(color.FgCyan, color.ReverseVideo),
color.New(color.FgWhite, color.ReverseVideo),
}
func main() {
fail := flag.Bool("fail", false, "terminate if any task fails with error")
flag.Parse()
tasks := flag.Args()
if len(tasks) > len(colors) {
log.Fatalf("Too many tasks! I only know %d colors.", len(colors))
}
wg := &sync.WaitGroup{}
wg.Add(len(tasks))
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
for i, command := range tasks {
colorOut := colorize(os.Stdout, i)
colorErr := colorize(os.Stderr, i)
fmt.Fprintf(colorErr, "Starting %q\n", command)
go func(i int, command string) {
if err := runCommand(ctx, command, colorOut, colorErr); err != nil {
fmt.Fprintf(colorErr, "command failed with %v\n", err)
if *fail { // terminate other tasks on failure
cancel()
}
}
colorOut.Close()
colorErr.Close()
wg.Done()
}(i, command)
}
wg.Wait()
}
func runCommand(ctx context.Context, command string, stdout, stderr io.Writer) error {
commandLine := append(shellCommand(), command)
cmd := exec.CommandContext(ctx, commandLine[0], commandLine[1:]...)
cmd.Stdout = stdout
cmd.Stderr = stderr
if err := cmd.Start(); err != nil {
return err
}
if err := cmd.Wait(); err != nil {
return err
}
return nil
}
func shellCommand() []string {
return []string{"bash", "-c"}
}
func colorize(dst io.Writer, i int) io.WriteCloser {
return &colorizer{
W: dst,
Prefix: fmt.Sprintf("%d> ", i),
Color: colors[i],
}
}
// mu protects writes to stdout and stderr
var mu sync.Mutex
type colorizer struct {
W io.Writer
Prefix string
Color *color.Color
trailer []byte
}
func (c *colorizer) write(prev, line []byte) (err error) {
mu.Lock()
_, err = c.Color.Fprintf(c.W, "%s%s%s\n", c.Prefix, prev, line)
mu.Unlock()
return
}
// Write writes the contents of p into W with color coding.
// Each Stream is output with an unique color.
// If p does not end with a newline, the trailing partial line
// is buffered and will be output on next write or on Close.
func (c *colorizer) Write(p []byte) (n int, err error) {
n = len(p)
for {
pos := bytes.IndexByte(p, '\n')
if pos == -1 { // incomplete last line
c.trailer = append(c.trailer[:0], p...)
return
}
line := p[:pos]
if err := c.write(c.trailer, line); err != nil {
return n, err
}
p = p[pos+1:]
c.trailer = nil
}
}
// Close writes trailing data not terminated with a newline.
func (c *colorizer) Close() error {
if len(c.trailer) > 0 {
return c.write(c.trailer, nil)
}
return nil
}