1package colorprofile
2
3import (
4 "bytes"
5 "fmt"
6 "image/color"
7 "io"
8 "strconv"
9
10 "github.com/charmbracelet/x/ansi"
11)
12
13// NewWriter creates a new color profile writer that downgrades color sequences
14// based on the detected color profile.
15//
16// If environ is nil, it will use os.Environ() to get the environment variables.
17//
18// It queries the given writer to determine if it supports ANSI escape codes.
19// If it does, along with the given environment variables, it will determine
20// the appropriate color profile to use for color formatting.
21//
22// This respects the NO_COLOR, CLICOLOR, and CLICOLOR_FORCE environment variables.
23func NewWriter(w io.Writer, environ []string) *Writer {
24 return &Writer{
25 Forward: w,
26 Profile: Detect(w, environ),
27 }
28}
29
30// Writer represents a color profile writer that writes ANSI sequences to the
31// underlying writer.
32type Writer struct {
33 Forward io.Writer
34 Profile Profile
35}
36
37// Write writes the given text to the underlying writer.
38func (w *Writer) Write(p []byte) (int, error) {
39 switch w.Profile {
40 case TrueColor:
41 return w.Forward.Write(p) //nolint:wrapcheck
42 case NoTTY:
43 return io.WriteString(w.Forward, ansi.Strip(string(p))) //nolint:wrapcheck
44 case Ascii, ANSI, ANSI256:
45 return w.downsample(p)
46 default:
47 return 0, fmt.Errorf("invalid profile: %v", w.Profile)
48 }
49}
50
51// downsample downgrades the given text to the appropriate color profile.
52func (w *Writer) downsample(p []byte) (int, error) {
53 var buf bytes.Buffer
54 var state byte
55
56 parser := ansi.GetParser()
57 defer ansi.PutParser(parser)
58
59 for len(p) > 0 {
60 parser.Reset()
61 seq, _, read, newState := ansi.DecodeSequence(p, state, parser)
62
63 switch {
64 case ansi.HasCsiPrefix(seq) && parser.Command() == 'm':
65 handleSgr(w, parser, &buf)
66 default:
67 // If we're not a style SGR sequence, just write the bytes.
68 if n, err := buf.Write(seq); err != nil {
69 return n, err //nolint:wrapcheck
70 }
71 }
72
73 p = p[read:]
74 state = newState
75 }
76
77 return w.Forward.Write(buf.Bytes()) //nolint:wrapcheck
78}
79
80// WriteString writes the given text to the underlying writer.
81func (w *Writer) WriteString(s string) (n int, err error) {
82 return w.Write([]byte(s))
83}
84
85func handleSgr(w *Writer, p *ansi.Parser, buf *bytes.Buffer) {
86 var style ansi.Style
87 params := p.Params()
88 for i := 0; i < len(params); i++ {
89 param := params[i]
90
91 switch param := param.Param(0); param {
92 case 0:
93 // SGR default parameter is 0. We use an empty string to reduce the
94 // number of bytes written to the buffer.
95 style = append(style, "")
96 case 30, 31, 32, 33, 34, 35, 36, 37: // 8-bit foreground color
97 if w.Profile < ANSI {
98 continue
99 }
100 style = style.ForegroundColor(
101 w.Profile.Convert(ansi.BasicColor(param - 30))) //nolint:gosec
102 case 38: // 16 or 24-bit foreground color
103 var c color.Color
104 if n := ansi.ReadStyleColor(params[i:], &c); n > 0 {
105 i += n - 1
106 }
107 if w.Profile < ANSI {
108 continue
109 }
110 style = style.ForegroundColor(w.Profile.Convert(c))
111 case 39: // default foreground color
112 if w.Profile < ANSI {
113 continue
114 }
115 style = style.DefaultForegroundColor()
116 case 40, 41, 42, 43, 44, 45, 46, 47: // 8-bit background color
117 if w.Profile < ANSI {
118 continue
119 }
120 style = style.BackgroundColor(
121 w.Profile.Convert(ansi.BasicColor(param - 40))) //nolint:gosec
122 case 48: // 16 or 24-bit background color
123 var c color.Color
124 if n := ansi.ReadStyleColor(params[i:], &c); n > 0 {
125 i += n - 1
126 }
127 if w.Profile < ANSI {
128 continue
129 }
130 style = style.BackgroundColor(w.Profile.Convert(c))
131 case 49: // default background color
132 if w.Profile < ANSI {
133 continue
134 }
135 style = style.DefaultBackgroundColor()
136 case 58: // 16 or 24-bit underline color
137 var c color.Color
138 if n := ansi.ReadStyleColor(params[i:], &c); n > 0 {
139 i += n - 1
140 }
141 if w.Profile < ANSI {
142 continue
143 }
144 style = style.UnderlineColor(w.Profile.Convert(c))
145 case 59: // default underline color
146 if w.Profile < ANSI {
147 continue
148 }
149 style = style.DefaultUnderlineColor()
150 case 90, 91, 92, 93, 94, 95, 96, 97: // 8-bit bright foreground color
151 if w.Profile < ANSI {
152 continue
153 }
154 style = style.ForegroundColor(
155 w.Profile.Convert(ansi.BasicColor(param - 90 + 8))) //nolint:gosec
156 case 100, 101, 102, 103, 104, 105, 106, 107: // 8-bit bright background color
157 if w.Profile < ANSI {
158 continue
159 }
160 style = style.BackgroundColor(
161 w.Profile.Convert(ansi.BasicColor(param - 100 + 8))) //nolint:gosec
162 default:
163 // If this is not a color attribute, just append it to the style.
164 style = append(style, strconv.Itoa(param))
165 }
166 }
167
168 _, _ = buf.WriteString(style.String())
169}