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// Strip removes ANSI escape codes from a string.
12func Strip(s string) string {
13 var (
14 buf bytes.Buffer // buffer for collecting printable characters
15 ri int // rune index
16 rw int // rune width
17 pstate = parser.GroundState // initial state
18 )
19
20 // This implements a subset of the Parser to only collect runes and
21 // printable characters.
22 for i := range len(s) {
23 if pstate == parser.Utf8State {
24 // During this state, collect rw bytes to form a valid rune in the
25 // buffer. After getting all the rune bytes into the buffer,
26 // transition to GroundState and reset the counters.
27 buf.WriteByte(s[i])
28 ri++
29 if ri < rw {
30 continue
31 }
32 pstate = parser.GroundState
33 ri = 0
34 rw = 0
35 continue
36 }
37
38 state, action := parser.Table.Transition(pstate, s[i])
39 switch action {
40 case parser.CollectAction:
41 if state == parser.Utf8State {
42 // This action happens when we transition to the Utf8State.
43 rw = utf8ByteLen(s[i])
44 buf.WriteByte(s[i])
45 ri++
46 }
47 case parser.PrintAction, parser.ExecuteAction:
48 // collects printable ASCII and non-printable characters
49 buf.WriteByte(s[i])
50 }
51
52 // Transition to the next state.
53 // The Utf8State is managed separately above.
54 if pstate != parser.Utf8State {
55 pstate = state
56 }
57 }
58
59 return buf.String()
60}
61
62// StringWidth returns the width of a string in cells. This is the number of
63// cells that the string will occupy when printed in a terminal. ANSI escape
64// codes are ignored and wide characters (such as East Asians and emojis) are
65// accounted for.
66// This treats the text as a sequence of grapheme clusters.
67func StringWidth(s string) int {
68 return stringWidth(GraphemeWidth, s)
69}
70
71// StringWidthWc returns the width of a string in cells. This is the number of
72// cells that the string will occupy when printed in a terminal. ANSI escape
73// codes are ignored and wide characters (such as East Asians and emojis) are
74// accounted for.
75// This treats the text as a sequence of wide characters and runes.
76func StringWidthWc(s string) int {
77 return stringWidth(WcWidth, s)
78}
79
80func stringWidth(m Method, s string) int {
81 if s == "" {
82 return 0
83 }
84
85 var (
86 pstate = parser.GroundState // initial state
87 cluster string
88 width int
89 )
90
91 for i := 0; i < len(s); i++ {
92 state, action := parser.Table.Transition(pstate, s[i])
93 if state == parser.Utf8State {
94 var w int
95 cluster, _, w, _ = uniseg.FirstGraphemeClusterInString(s[i:], -1)
96 if m == WcWidth {
97 w = runewidth.StringWidth(cluster)
98 }
99 width += w
100 i += len(cluster) - 1
101 pstate = parser.GroundState
102 continue
103 }
104
105 if action == parser.PrintAction {
106 width++
107 }
108
109 pstate = state
110 }
111
112 return width
113}