ansi.go

Documentation: github.com/rivo/tview

     1  package tview
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"strconv"
     8  	"strings"
     9  )
    10  
    11  // The states of the ANSI escape code parser.
    12  const (
    13  	ansiText = iota
    14  	ansiEscape
    15  	ansiSubstring
    16  	ansiControlSequence
    17  )
    18  
    19  // ansi is a io.Writer which translates ANSI escape codes into tview color
    20  // tags.
    21  type ansi struct {
    22  	io.Writer
    23  
    24  	// Reusable buffers.
    25  	buffer                        *bytes.Buffer // The entire output text of one Write().
    26  	csiParameter, csiIntermediate *bytes.Buffer // Partial CSI strings.
    27  	attributes                    string        // The buffer's current text attributes (a tview attribute string).
    28  
    29  	// The current state of the parser. One of the ansi constants.
    30  	state int
    31  }
    32  
    33  // ANSIWriter returns an io.Writer which translates any ANSI escape codes
    34  // written to it into tview color tags. Other escape codes don't have an effect
    35  // and are simply removed. The translated text is written to the provided
    36  // writer.
    37  func ANSIWriter(writer io.Writer) io.Writer {
    38  	return &ansi{
    39  		Writer:          writer,
    40  		buffer:          new(bytes.Buffer),
    41  		csiParameter:    new(bytes.Buffer),
    42  		csiIntermediate: new(bytes.Buffer),
    43  		state:           ansiText,
    44  	}
    45  }
    46  
    47  // Write parses the given text as a string of runes, translates ANSI escape
    48  // codes to color tags and writes them to the output writer.
    49  func (a *ansi) Write(text []byte) (int, error) {
    50  	defer func() {
    51  		a.buffer.Reset()
    52  	}()
    53  
    54  	for _, r := range string(text) {
    55  		switch a.state {
    56  
    57  		// We just entered an escape sequence.
    58  		case ansiEscape:
    59  			switch r {
    60  			case '[': // Control Sequence Introducer.
    61  				a.csiParameter.Reset()
    62  				a.csiIntermediate.Reset()
    63  				a.state = ansiControlSequence
    64  			case 'c': // Reset.
    65  				fmt.Fprint(a.buffer, "[-:-:-]")
    66  				a.state = ansiText
    67  			case 'P', ']', 'X', '^', '_': // Substrings and commands.
    68  				a.state = ansiSubstring
    69  			default: // Ignore.
    70  				a.state = ansiText
    71  			}
    72  
    73  		// CSI Sequences.
    74  		case ansiControlSequence:
    75  			switch {
    76  			case r >= 0x30 && r <= 0x3f: // Parameter bytes.
    77  				if _, err := a.csiParameter.WriteRune(r); err != nil {
    78  					return 0, err
    79  				}
    80  			case r >= 0x20 && r <= 0x2f: // Intermediate bytes.
    81  				if _, err := a.csiIntermediate.WriteRune(r); err != nil {
    82  					return 0, err
    83  				}
    84  			case r >= 0x40 && r <= 0x7e: // Final byte.
    85  				switch r {
    86  				case 'E': // Next line.
    87  					count, _ := strconv.Atoi(a.csiParameter.String())
    88  					if count == 0 {
    89  						count = 1
    90  					}
    91  					fmt.Fprint(a.buffer, strings.Repeat("\n", count))
    92  				case 'm': // Select Graphic Rendition.
    93  					var background, foreground string
    94  					params := a.csiParameter.String()
    95  					fields := strings.Split(params, ";")
    96  					if len(params) == 0 || len(fields) == 1 && fields[0] == "0" {
    97  						// Reset.
    98  						a.attributes = ""
    99  						if _, err := a.buffer.WriteString("[-:-:-]"); err != nil {
   100  							return 0, err
   101  						}
   102  						break
   103  					}
   104  					lookupColor := func(colorNumber int) string {
   105  						if colorNumber < 0 || colorNumber > 15 {
   106  							return "black"
   107  						}
   108  						return []string{
   109  							"black",
   110  							"maroon",
   111  							"green",
   112  							"olive",
   113  							"navy",
   114  							"purple",
   115  							"teal",
   116  							"silver",
   117  							"gray",
   118  							"red",
   119  							"lime",
   120  							"yellow",
   121  							"blue",
   122  							"fuchsia",
   123  							"aqua",
   124  							"white",
   125  						}[colorNumber]
   126  					}
   127  				FieldLoop:
   128  					for index, field := range fields {
   129  						switch field {
   130  						case "1", "01":
   131  							if !strings.ContainsRune(a.attributes, 'b') {
   132  								a.attributes += "b"
   133  							}
   134  						case "2", "02":
   135  							if !strings.ContainsRune(a.attributes, 'd') {
   136  								a.attributes += "d"
   137  							}
   138  						case "4", "04":
   139  							if !strings.ContainsRune(a.attributes, 'u') {
   140  								a.attributes += "u"
   141  							}
   142  						case "5", "05":
   143  							if !strings.ContainsRune(a.attributes, 'l') {
   144  								a.attributes += "l"
   145  							}
   146  						case "22":
   147  							if i := strings.IndexRune(a.attributes, 'b'); i >= 0 {
   148  								a.attributes = a.attributes[:i] + a.attributes[i+1:]
   149  							}
   150  							if i := strings.IndexRune(a.attributes, 'd'); i >= 0 {
   151  								a.attributes = a.attributes[:i] + a.attributes[i+1:]
   152  							}
   153  						case "24":
   154  							if i := strings.IndexRune(a.attributes, 'u'); i >= 0 {
   155  								a.attributes = a.attributes[:i] + a.attributes[i+1:]
   156  							}
   157  						case "25":
   158  							if i := strings.IndexRune(a.attributes, 'l'); i >= 0 {
   159  								a.attributes = a.attributes[:i] + a.attributes[i+1:]
   160  							}
   161  						case "30", "31", "32", "33", "34", "35", "36", "37":
   162  							colorNumber, _ := strconv.Atoi(field)
   163  							foreground = lookupColor(colorNumber - 30)
   164  						case "39":
   165  							foreground = "-"
   166  						case "40", "41", "42", "43", "44", "45", "46", "47":
   167  							colorNumber, _ := strconv.Atoi(field)
   168  							background = lookupColor(colorNumber - 40)
   169  						case "49":
   170  							background = "-"
   171  						case "90", "91", "92", "93", "94", "95", "96", "97":
   172  							colorNumber, _ := strconv.Atoi(field)
   173  							foreground = lookupColor(colorNumber - 82)
   174  						case "100", "101", "102", "103", "104", "105", "106", "107":
   175  							colorNumber, _ := strconv.Atoi(field)
   176  							background = lookupColor(colorNumber - 92)
   177  						case "38", "48":
   178  							var color string
   179  							if len(fields) > index+1 {
   180  								if fields[index+1] == "5" && len(fields) > index+2 { // 8-bit colors.
   181  									colorNumber, _ := strconv.Atoi(fields[index+2])
   182  									if colorNumber <= 15 {
   183  										color = lookupColor(colorNumber)
   184  									} else if colorNumber <= 231 {
   185  										red := (colorNumber - 16) / 36
   186  										green := ((colorNumber - 16) / 6) % 6
   187  										blue := (colorNumber - 16) % 6
   188  										color = fmt.Sprintf("#%02x%02x%02x", 255*red/5, 255*green/5, 255*blue/5)
   189  									} else if colorNumber <= 255 {
   190  										grey := 255 * (colorNumber - 232) / 23
   191  										color = fmt.Sprintf("#%02x%02x%02x", grey, grey, grey)
   192  									}
   193  								} else if fields[index+1] == "2" && len(fields) > index+4 { // 24-bit colors.
   194  									red, _ := strconv.Atoi(fields[index+2])
   195  									green, _ := strconv.Atoi(fields[index+3])
   196  									blue, _ := strconv.Atoi(fields[index+4])
   197  									color = fmt.Sprintf("#%02x%02x%02x", red, green, blue)
   198  								}
   199  							}
   200  							if len(color) > 0 {
   201  								if field == "38" {
   202  									foreground = color
   203  								} else {
   204  									background = color
   205  								}
   206  							}
   207  							break FieldLoop
   208  						}
   209  					}
   210  					var colon string
   211  					if len(a.attributes) > 0 {
   212  						colon = ":"
   213  					}
   214  					if len(foreground) > 0 || len(background) > 0 || len(a.attributes) > 0 {
   215  						fmt.Fprintf(a.buffer, "[%s:%s%s%s]", foreground, background, colon, a.attributes)
   216  					}
   217  				}
   218  				a.state = ansiText
   219  			default: // Undefined byte.
   220  				a.state = ansiText // Abort CSI.
   221  			}
   222  
   223  			// We just entered a substring/command sequence.
   224  		case ansiSubstring:
   225  			if r == 27 { // Most likely the end of the substring.
   226  				a.state = ansiEscape
   227  			} // Ignore all other characters.
   228  
   229  			// "ansiText" and all others.
   230  		default:
   231  			if r == 27 {
   232  				// This is the start of an escape sequence.
   233  				a.state = ansiEscape
   234  			} else {
   235  				// Just a regular rune. Send to buffer.
   236  				if _, err := a.buffer.WriteRune(r); err != nil {
   237  					return 0, err
   238  				}
   239  			}
   240  		}
   241  	}
   242  
   243  	// Write buffer to target writer.
   244  	n, err := a.buffer.WriteTo(a.Writer)
   245  	if err != nil {
   246  		return int(n), err
   247  	}
   248  	return len(text), nil
   249  }
   250  
   251  // TranslateANSI replaces ANSI escape sequences found in the provided string
   252  // with tview's color tags and returns the resulting string.
   253  func TranslateANSI(text string) string {
   254  	var buffer bytes.Buffer
   255  	writer := ANSIWriter(&buffer)
   256  	writer.Write([]byte(text))
   257  	return buffer.String()
   258  }
   259
View as plain text