github.com/CUCUMBER/godog@v0.7.9/color_tag_test.go (about) 1 package godog 2 3 import ( 4 "bytes" 5 "fmt" 6 "io" 7 "strings" 8 "testing" 9 10 "github.com/DATA-DOG/godog/colors" 11 ) 12 13 type csiState int 14 15 const ( 16 outsideCsiCode csiState = iota 17 firstCsiCode 18 secondCsiCode 19 ) 20 21 type tagColorWriter struct { 22 w io.Writer 23 state csiState 24 paramStartBuf bytes.Buffer 25 paramBuf bytes.Buffer 26 tag string 27 } 28 29 const ( 30 firstCsiChar byte = '\x1b' 31 secondeCsiChar byte = '[' 32 separatorChar byte = ';' 33 sgrCode byte = 'm' 34 ) 35 36 const ( 37 ansiReset = "0" 38 ansiIntensityOn = "1" 39 ansiIntensityOff = "21" 40 ansiUnderlineOn = "4" 41 ansiUnderlineOff = "24" 42 ansiBlinkOn = "5" 43 ansiBlinkOff = "25" 44 45 ansiForegroundBlack = "30" 46 ansiForegroundRed = "31" 47 ansiForegroundGreen = "32" 48 ansiForegroundYellow = "33" 49 ansiForegroundBlue = "34" 50 ansiForegroundMagenta = "35" 51 ansiForegroundCyan = "36" 52 ansiForegroundWhite = "37" 53 ansiForegroundDefault = "39" 54 55 ansiBackgroundBlack = "40" 56 ansiBackgroundRed = "41" 57 ansiBackgroundGreen = "42" 58 ansiBackgroundYellow = "43" 59 ansiBackgroundBlue = "44" 60 ansiBackgroundMagenta = "45" 61 ansiBackgroundCyan = "46" 62 ansiBackgroundWhite = "47" 63 ansiBackgroundDefault = "49" 64 65 ansiLightForegroundGray = "90" 66 ansiLightForegroundRed = "91" 67 ansiLightForegroundGreen = "92" 68 ansiLightForegroundYellow = "93" 69 ansiLightForegroundBlue = "94" 70 ansiLightForegroundMagenta = "95" 71 ansiLightForegroundCyan = "96" 72 ansiLightForegroundWhite = "97" 73 74 ansiLightBackgroundGray = "100" 75 ansiLightBackgroundRed = "101" 76 ansiLightBackgroundGreen = "102" 77 ansiLightBackgroundYellow = "103" 78 ansiLightBackgroundBlue = "104" 79 ansiLightBackgroundMagenta = "105" 80 ansiLightBackgroundCyan = "106" 81 ansiLightBackgroundWhite = "107" 82 ) 83 84 var colorMap = map[string]string{ 85 ansiForegroundBlack: "black", 86 ansiForegroundRed: "red", 87 ansiForegroundGreen: "green", 88 ansiForegroundYellow: "yellow", 89 ansiForegroundBlue: "blue", 90 ansiForegroundMagenta: "magenta", 91 ansiForegroundCyan: "cyan", 92 ansiForegroundWhite: "white", 93 ansiForegroundDefault: "", 94 } 95 96 func (cw *tagColorWriter) flushBuffer() (int, error) { 97 return cw.flushTo(cw.w) 98 } 99 100 func (cw *tagColorWriter) resetBuffer() (int, error) { 101 return cw.flushTo(nil) 102 } 103 104 func (cw *tagColorWriter) flushTo(w io.Writer) (int, error) { 105 var n1, n2 int 106 var err error 107 108 startBytes := cw.paramStartBuf.Bytes() 109 cw.paramStartBuf.Reset() 110 if w != nil { 111 n1, err = cw.w.Write(startBytes) 112 if err != nil { 113 return n1, err 114 } 115 } else { 116 n1 = len(startBytes) 117 } 118 paramBytes := cw.paramBuf.Bytes() 119 cw.paramBuf.Reset() 120 if w != nil { 121 n2, err = cw.w.Write(paramBytes) 122 if err != nil { 123 return n1 + n2, err 124 } 125 } else { 126 n2 = len(paramBytes) 127 } 128 return n1 + n2, nil 129 } 130 131 func isParameterChar(b byte) bool { 132 return ('0' <= b && b <= '9') || b == separatorChar 133 } 134 135 func (cw *tagColorWriter) Write(p []byte) (int, error) { 136 r, nw, first, last := 0, 0, 0, 0 137 138 var err error 139 for i, ch := range p { 140 switch cw.state { 141 case outsideCsiCode: 142 if ch == firstCsiChar { 143 cw.paramStartBuf.WriteByte(ch) 144 cw.state = firstCsiCode 145 } 146 case firstCsiCode: 147 switch ch { 148 case firstCsiChar: 149 cw.paramStartBuf.WriteByte(ch) 150 break 151 case secondeCsiChar: 152 cw.paramStartBuf.WriteByte(ch) 153 cw.state = secondCsiCode 154 last = i - 1 155 default: 156 cw.resetBuffer() 157 cw.state = outsideCsiCode 158 } 159 case secondCsiCode: 160 if isParameterChar(ch) { 161 cw.paramBuf.WriteByte(ch) 162 } else { 163 nw, err = cw.w.Write(p[first:last]) 164 r += nw 165 if err != nil { 166 return r, err 167 } 168 first = i + 1 169 if ch == sgrCode { 170 cw.changeColor() 171 } 172 n, _ := cw.resetBuffer() 173 // Add one more to the size of the buffer for the last ch 174 r += n + 1 175 176 cw.state = outsideCsiCode 177 } 178 default: 179 cw.state = outsideCsiCode 180 } 181 } 182 183 if cw.state == outsideCsiCode { 184 nw, err = cw.w.Write(p[first:len(p)]) 185 r += nw 186 } 187 188 return r, err 189 } 190 191 func (cw *tagColorWriter) changeColor() { 192 strParam := cw.paramBuf.String() 193 if len(strParam) <= 0 { 194 strParam = "0" 195 } 196 csiParam := strings.Split(strParam, string(separatorChar)) 197 for _, p := range csiParam { 198 c, ok := colorMap[p] 199 switch { 200 case !ok: 201 switch p { 202 case ansiReset: 203 fmt.Fprint(cw.w, "</"+cw.tag+">") 204 cw.tag = "" 205 case ansiIntensityOn: 206 cw.tag = "bold-" + cw.tag 207 case ansiIntensityOff: 208 case ansiUnderlineOn: 209 case ansiUnderlineOff: 210 case ansiBlinkOn: 211 case ansiBlinkOff: 212 default: 213 // unknown code 214 } 215 default: 216 cw.tag += c 217 fmt.Fprint(cw.w, "<"+cw.tag+">") 218 } 219 } 220 } 221 222 func TestTagColorWriter(t *testing.T) { 223 var buf bytes.Buffer 224 w := &tagColorWriter{w: &buf} 225 226 s := fmt.Sprintf("text %s then %s", colors.Red("in red"), colors.Yellow("yel")) 227 fmt.Fprint(w, s) 228 229 expected := "text <red>in red</red> then <yellow>yel</yellow>" 230 if buf.String() != expected { 231 t.Fatalf("expected `%s` but got `%s`", expected, buf.String()) 232 } 233 } 234 235 func TestTagBoldColorWriter(t *testing.T) { 236 var buf bytes.Buffer 237 w := &tagColorWriter{w: &buf} 238 239 s := fmt.Sprintf( 240 "text %s then %s", 241 colors.Bold(colors.Red)("in red"), 242 colors.Bold(colors.Yellow)("yel"), 243 ) 244 fmt.Fprint(w, s) 245 246 expected := "text <bold-red>in red</bold-red> then <bold-yellow>yel</bold-yellow>" 247 if buf.String() != expected { 248 t.Fatalf("expected `%s` but got `%s`", expected, buf.String()) 249 } 250 }