1
2 package cmdlog
3
4 import (
5 "bytes"
6 "io"
7 "log"
8 "os"
9 "sync"
10 "sync/atomic"
11 "testing"
12 "time"
13
14 "oss.terrastruct.com/util-go/xos"
15 "oss.terrastruct.com/util-go/xterm"
16 )
17
18 var timeNow = time.Now
19
20 const defaultTSFormat = "15:04:05"
21
22 func init() {
23 l := New(xos.NewEnv(os.Environ()), os.Stderr)
24 l.SetTS(true)
25 l = l.WithPrefix(xterm.Blue, "stdlog")
26
27 log.SetOutput(l.NoLevel.Writer())
28 log.SetPrefix(l.NoLevel.Prefix())
29 log.SetFlags(l.NoLevel.Flags())
30 }
31
32 type Logger struct {
33 env *xos.Env
34 w io.Writer
35 tsw *tsWriter
36 dw *debugWriter
37
38 NoLevel *log.Logger
39 Debug *log.Logger
40 Success *log.Logger
41 Info *log.Logger
42 Warn *log.Logger
43 Error *log.Logger
44 }
45
46 func (l *Logger) GetTS() bool {
47 l.tsw.mu.Lock()
48 defer l.tsw.mu.Unlock()
49 return l.tsw.enabled
50 }
51
52 func (l *Logger) GetTSFormat() string {
53 l.tsw.mu.Lock()
54 defer l.tsw.mu.Unlock()
55 return l.tsw.tsfmt
56 }
57
58 func (l *Logger) GetDebug() bool {
59 return l.dw.debug()
60 }
61
62 func (l *Logger) SetTS(enabled bool) {
63 l.tsw.mu.Lock()
64 l.tsw.enabled = enabled
65 l.tsw.mu.Unlock()
66 }
67
68 func (l *Logger) SetTSFormat(tsfmt string) {
69 l.tsw.mu.Lock()
70 l.tsw.tsfmt = tsfmt
71 l.tsw.mu.Unlock()
72 }
73
74 func (l *Logger) SetDebug(enabled bool) {
75 vi := int64(0)
76 if enabled {
77 vi = 1
78 }
79 atomic.StoreInt64(&l.dw.flag, vi)
80 }
81
82 func New(env *xos.Env, w io.Writer) *Logger {
83 tsw := &tsWriter{w: w, tsfmt: defaultTSFormat}
84 dw := &debugWriter{w: tsw, env: env}
85 l := &Logger{
86 env: env,
87 w: w,
88 dw: dw,
89 tsw: tsw,
90 }
91 l.init("")
92 return l
93 }
94
95 func (l *Logger) init(prefix string) {
96 l.NoLevel = log.New(prefixWriter{l.tsw, prefix}, "", 0)
97
98 if prefix != "" {
99 prefix += " "
100 }
101 l.Debug = log.New(prefixWriter{l.dw, prefix + xterm.Prefix(l.env, l.w, "", "debug")}, "", 0)
102 l.Success = log.New(prefixWriter{l.tsw, prefix + xterm.Prefix(l.env, l.w, xterm.Green, "success")}, "", 0)
103 l.Info = log.New(prefixWriter{l.tsw, prefix + xterm.Prefix(l.env, l.w, xterm.Blue, "info")}, "", 0)
104 l.Warn = log.New(prefixWriter{l.tsw, prefix + xterm.Prefix(l.env, l.w, xterm.Yellow, "warn")}, "", 0)
105 l.Error = log.New(prefixWriter{l.tsw, prefix + xterm.Prefix(l.env, l.w, xterm.Red, "err")}, "", 0)
106 }
107
108 type prefixWriter struct {
109 w io.Writer
110 prefix string
111 }
112
113 func (pw prefixWriter) Write(p []byte) (int, error) {
114 lines := bytes.Split(p, []byte("\n"))
115 p2 := make([]byte, 0, (len(pw.prefix)+1)*len(lines)+len(p))
116
117 for _, l := range lines[:len(lines)-1] {
118 prefix := pw.prefix
119 if len(l) > 0 {
120 prefix += " "
121 }
122 p2 = append(p2, prefix...)
123 p2 = append(p2, l...)
124 p2 = append(p2, '\n')
125 }
126
127 n, err := pw.w.Write(p2)
128 if n > len(p) {
129 n = len(p)
130 }
131 return n, err
132 }
133
134 type debugWriter struct {
135 w io.Writer
136 flag int64
137 env *xos.Env
138 }
139
140 func (dw *debugWriter) debug() bool {
141 if atomic.LoadInt64(&dw.flag) == 0 {
142 return dw.env.Debug()
143 }
144 return true
145 }
146
147 func (dw *debugWriter) Write(p []byte) (int, error) {
148 if !dw.debug() {
149 return len(p), nil
150 }
151 return dw.w.Write(p)
152 }
153
154 type tsWriter struct {
155 w io.Writer
156
157 mu sync.Mutex
158 tsfmt string
159 enabled bool
160 }
161
162 func (tsw *tsWriter) Write(p []byte) (int, error) {
163 tsw.mu.Lock()
164 enabled := tsw.enabled
165 tsfmt := tsw.tsfmt
166 tsw.mu.Unlock()
167
168 if !enabled {
169 return tsw.w.Write(p)
170 }
171
172 ts := timeNow().Format(tsfmt)
173 prefix := []byte("[" + ts + "]")
174
175 lines := bytes.Split(p, []byte("\n"))
176 p2 := make([]byte, 0, (len(prefix)+1)*len(lines)+len(p))
177
178 for _, l := range lines[:len(lines)-1] {
179 prefix := prefix
180 if len(l) > 0 {
181 prefix = append(prefix, ' ')
182 }
183 p2 = append(p2, prefix...)
184 p2 = append(p2, l...)
185 p2 = append(p2, '\n')
186 }
187
188 n, err := tsw.w.Write(p2)
189 if n > len(p) {
190 n = len(p)
191 }
192 return n, err
193 }
194
195 func NewTB(env *xos.Env, tb testing.TB) *Logger {
196 return New(env, tbWriter{tb})
197 }
198
199 type tbWriter struct {
200 tb testing.TB
201 }
202
203 func (tbw tbWriter) Write(p []byte) (int, error) {
204 tbw.tb.Logf("%s", p)
205 return len(p), nil
206 }
207
208
209 func (tbWriter) Fd() uintptr {
210 return os.Stderr.Fd()
211 }
212
213 func (l *Logger) WithCCPrefix(s string) *Logger {
214 return l.withPrefix(xterm.CCPrefix(l.env, l.w, s))
215 }
216
217 func (l *Logger) WithPrefix(caps, s string) *Logger {
218 return l.withPrefix(xterm.Prefix(l.env, l.w, caps, s))
219 }
220
221 func (l *Logger) withPrefix(s string) *Logger {
222 l2 := new(Logger)
223 *l2 = *l
224
225 prefix := l.NoLevel.Writer().(prefixWriter).prefix
226 if len(s) > 0 {
227 if len(prefix) > 0 {
228 prefix += " "
229 }
230 prefix += s
231 }
232 l2.init(prefix)
233 return l2
234 }
235
View as plain text