main
1package ansi
2
3import (
4 "bytes"
5
6 "github.com/charmbracelet/x/ansi/parser"
7 "github.com/mattn/go-runewidth"
8 "github.com/rivo/uniseg"
9)
10
11// Cut the string, without adding any prefix or tail strings. This function is
12// aware of ANSI escape codes and will not break them, and accounts for
13// wide-characters (such as East-Asian characters and emojis). Note that the
14// [left] parameter is inclusive, while [right] isn't.
15// This treats the text as a sequence of graphemes.
16func Cut(s string, left, right int) string {
17 return cut(GraphemeWidth, s, left, right)
18}
19
20// CutWc the string, without adding any prefix or tail strings. This function is
21// aware of ANSI escape codes and will not break them, and accounts for
22// wide-characters (such as East-Asian characters and emojis). Note that the
23// [left] parameter is inclusive, while [right] isn't.
24// This treats the text as a sequence of wide characters and runes.
25func CutWc(s string, left, right int) string {
26 return cut(WcWidth, s, left, right)
27}
28
29func cut(m Method, s string, left, right int) string {
30 if right <= left {
31 return ""
32 }
33
34 truncate := Truncate
35 truncateLeft := TruncateLeft
36 if m == WcWidth {
37 truncate = TruncateWc
38 truncateLeft = TruncateWc
39 }
40
41 if left == 0 {
42 return truncate(s, right, "")
43 }
44 return truncateLeft(Truncate(s, right, ""), left, "")
45}
46
47// Truncate truncates a string to a given length, adding a tail to the end if
48// the string is longer than the given length. This function is aware of ANSI
49// escape codes and will not break them, and accounts for wide-characters (such
50// as East-Asian characters and emojis).
51// This treats the text as a sequence of graphemes.
52func Truncate(s string, length int, tail string) string {
53 return truncate(GraphemeWidth, s, length, tail)
54}
55
56// TruncateWc truncates a string to a given length, adding a tail to the end if
57// the string is longer than the given length. This function is aware of ANSI
58// escape codes and will not break them, and accounts for wide-characters (such
59// as East-Asian characters and emojis).
60// This treats the text as a sequence of wide characters and runes.
61func TruncateWc(s string, length int, tail string) string {
62 return truncate(WcWidth, s, length, tail)
63}
64
65func truncate(m Method, s string, length int, tail string) string {
66 if sw := StringWidth(s); sw <= length {
67 return s
68 }
69
70 tw := StringWidth(tail)
71 length -= tw
72 if length < 0 {
73 return ""
74 }
75
76 var cluster []byte
77 var buf bytes.Buffer
78 curWidth := 0
79 ignoring := false
80 pstate := parser.GroundState // initial state
81 b := []byte(s)
82 i := 0
83
84 // Here we iterate over the bytes of the string and collect printable
85 // characters and runes. We also keep track of the width of the string
86 // in cells.
87 //
88 // Once we reach the given length, we start ignoring characters and only
89 // collect ANSI escape codes until we reach the end of string.
90 for i < len(b) {
91 state, action := parser.Table.Transition(pstate, b[i])
92 if state == parser.Utf8State {
93 // This action happens when we transition to the Utf8State.
94 var width int
95 cluster, _, width, _ = uniseg.FirstGraphemeCluster(b[i:], -1)
96 if m == WcWidth {
97 width = runewidth.StringWidth(string(cluster))
98 }
99
100 // increment the index by the length of the cluster
101 i += len(cluster)
102
103 // Are we ignoring? Skip to the next byte
104 if ignoring {
105 continue
106 }
107
108 // Is this gonna be too wide?
109 // If so write the tail and stop collecting.
110 if curWidth+width > length && !ignoring {
111 ignoring = true
112 buf.WriteString(tail)
113 }
114
115 if curWidth+width > length {
116 continue
117 }
118
119 curWidth += width
120 buf.Write(cluster)
121
122 // Done collecting, now we're back in the ground state.
123 pstate = parser.GroundState
124 continue
125 }
126
127 switch action {
128 case parser.PrintAction:
129 // Is this gonna be too wide?
130 // If so write the tail and stop collecting.
131 if curWidth >= length && !ignoring {
132 ignoring = true
133 buf.WriteString(tail)
134 }
135
136 // Skip to the next byte if we're ignoring
137 if ignoring {
138 i++
139 continue
140 }
141
142 // collects printable ASCII
143 curWidth++
144 fallthrough
145 default:
146 buf.WriteByte(b[i])
147 i++
148 }
149
150 // Transition to the next state.
151 pstate = state
152
153 // Once we reach the given length, we start ignoring runes and write
154 // the tail to the buffer.
155 if curWidth > length && !ignoring {
156 ignoring = true
157 buf.WriteString(tail)
158 }
159 }
160
161 return buf.String()
162}
163
164// TruncateLeft truncates a string from the left side by removing n characters,
165// adding a prefix to the beginning if the string is longer than n.
166// This function is aware of ANSI escape codes and will not break them, and
167// accounts for wide-characters (such as East-Asian characters and emojis).
168// This treats the text as a sequence of graphemes.
169func TruncateLeft(s string, n int, prefix string) string {
170 return truncateLeft(GraphemeWidth, s, n, prefix)
171}
172
173// TruncateLeftWc truncates a string from the left side by removing n characters,
174// adding a prefix to the beginning if the string is longer than n.
175// This function is aware of ANSI escape codes and will not break them, and
176// accounts for wide-characters (such as East-Asian characters and emojis).
177// This treats the text as a sequence of wide characters and runes.
178func TruncateLeftWc(s string, n int, prefix string) string {
179 return truncateLeft(WcWidth, s, n, prefix)
180}
181
182func truncateLeft(m Method, s string, n int, prefix string) string {
183 if n <= 0 {
184 return s
185 }
186
187 var cluster []byte
188 var buf bytes.Buffer
189 curWidth := 0
190 ignoring := true
191 pstate := parser.GroundState
192 b := []byte(s)
193 i := 0
194
195 for i < len(b) {
196 if !ignoring {
197 buf.Write(b[i:])
198 break
199 }
200
201 state, action := parser.Table.Transition(pstate, b[i])
202 if state == parser.Utf8State {
203 var width int
204 cluster, _, width, _ = uniseg.FirstGraphemeCluster(b[i:], -1)
205 if m == WcWidth {
206 width = runewidth.StringWidth(string(cluster))
207 }
208
209 i += len(cluster)
210 curWidth += width
211
212 if curWidth > n && ignoring {
213 ignoring = false
214 buf.WriteString(prefix)
215 }
216
217 if ignoring {
218 continue
219 }
220
221 if curWidth > n {
222 buf.Write(cluster)
223 }
224
225 pstate = parser.GroundState
226 continue
227 }
228
229 switch action {
230 case parser.PrintAction:
231 curWidth++
232
233 if curWidth > n && ignoring {
234 ignoring = false
235 buf.WriteString(prefix)
236 }
237
238 if ignoring {
239 i++
240 continue
241 }
242
243 fallthrough
244 default:
245 buf.WriteByte(b[i])
246 i++
247 }
248
249 pstate = state
250 if curWidth > n && ignoring {
251 ignoring = false
252 buf.WriteString(prefix)
253 }
254 }
255
256 return buf.String()
257}
258
259// ByteToGraphemeRange takes start and stop byte positions and converts them to
260// grapheme-aware char positions.
261// You can use this with [Truncate], [TruncateLeft], and [Cut].
262func ByteToGraphemeRange(str string, byteStart, byteStop int) (charStart, charStop int) {
263 bytePos, charPos := 0, 0
264 gr := uniseg.NewGraphemes(str)
265 for byteStart > bytePos {
266 if !gr.Next() {
267 break
268 }
269 bytePos += len(gr.Str())
270 charPos += max(1, gr.Width())
271 }
272 charStart = charPos
273 for byteStop > bytePos {
274 if !gr.Next() {
275 break
276 }
277 bytePos += len(gr.Str())
278 charPos += max(1, gr.Width())
279 }
280 charStop = charPos
281 return
282}