nt

ansi.go (8371B)

---

 package tview
 
 import (
 	"bytes"
 	"fmt"
 	"io"
 	"strconv"
 	"strings"
 )
 
 // The states of the ANSI escape code parser.
 const (
 	ansiText = iota
 	ansiEscape
 	ansiSubstring
 	ansiControlSequence
 )
 
 // ansi is a io.Writer which translates ANSI escape codes into tview color
 // tags.
 type ansi struct {
 	io.Writer
 
 	// Reusable buffers.
 	buffer                        *bytes.Buffer // The entire output text of one Write().
 	csiParameter, csiIntermediate *bytes.Buffer // Partial CSI strings.
 	attributes                    string        // The buffer's current text attributes (a tview attribute string).
 
 	// The current state of the parser. One of the ansi constants.
 	state int
 }
 
 // ANSIWriter returns an io.Writer which translates any ANSI escape codes
 // written to it into tview style tags. Other escape codes don't have an effect
 // and are simply removed. The translated text is written to the provided
 // writer.
 func ANSIWriter(writer io.Writer) io.Writer {
 	return &ansi{
 		Writer:          writer,
 		buffer:          new(bytes.Buffer),
 		csiParameter:    new(bytes.Buffer),
 		csiIntermediate: new(bytes.Buffer),
 		state:           ansiText,
 	}
 }
 
 // Write parses the given text as a string of runes, translates ANSI escape
 // codes to style tags and writes them to the output writer.
 func (a *ansi) Write(text []byte) (int, error) {
 	defer func() {
 		a.buffer.Reset()
 	}()
 
 	for _, r := range string(text) {
 		switch a.state {
 
 		// We just entered an escape sequence.
 		case ansiEscape:
 			switch r {
 			case '[': // Control Sequence Introducer.
 				a.csiParameter.Reset()
 				a.csiIntermediate.Reset()
 				a.state = ansiControlSequence
 			case 'c': // Reset.
 				fmt.Fprint(a.buffer, "[-:-:-]")
 				a.state = ansiText
 			case 'P', ']', 'X', '^', '_': // Substrings and commands.
 				a.state = ansiSubstring
 			default: // Ignore.
 				a.state = ansiText
 			}
 
 		// CSI Sequences.
 		case ansiControlSequence:
 			switch {
 			case r >= 0x30 && r <= 0x3f: // Parameter bytes.
 				if _, err := a.csiParameter.WriteRune(r); err != nil {
 					return 0, err
 				}
 			case r >= 0x20 && r <= 0x2f: // Intermediate bytes.
 				if _, err := a.csiIntermediate.WriteRune(r); err != nil {
 					return 0, err
 				}
 			case r >= 0x40 && r <= 0x7e: // Final byte.
 				switch r {
 				case 'E': // Next line.
 					count, _ := strconv.Atoi(a.csiParameter.String())
 					if count == 0 {
 						count = 1
 					}
 					fmt.Fprint(a.buffer, strings.Repeat("\n", count))
 				case 'm': // Select Graphic Rendition.
 					var background, foreground string
 					params := a.csiParameter.String()
 					fields := strings.Split(params, ";")
 					if len(params) == 0 || fields[0] == "" || fields[0] == "0" {
 						// Reset.
 						foreground = "-"
 						background = "-"
 						a.attributes = "-"
 					}
 					lookupColor := func(colorNumber int) string {
 						if colorNumber < 0 || colorNumber > 15 {
 							return "black"
 						}
 						return []string{
 							"black",
 							"maroon",
 							"green",
 							"olive",
 							"navy",
 							"purple",
 							"teal",
 							"silver",
 							"gray",
 							"red",
 							"lime",
 							"yellow",
 							"blue",
 							"fuchsia",
 							"aqua",
 							"white",
 						}[colorNumber]
 					}
 				FieldLoop:
 					for index, field := range fields {
 						switch field {
 						case "1", "01":
 							if !strings.ContainsRune(a.attributes, 'b') {
 								a.attributes += "b"
 							}
 						case "2", "02":
 							if !strings.ContainsRune(a.attributes, 'd') {
 								a.attributes += "d"
 							}
 						case "3", "03":
 							if !strings.ContainsRune(a.attributes, 'i') {
 								a.attributes += "i"
 							}
 						case "4", "04":
 							if !strings.ContainsRune(a.attributes, 'u') {
 								a.attributes += "u"
 							}
 						case "5", "05":
 							if !strings.ContainsRune(a.attributes, 'l') {
 								a.attributes += "l"
 							}
 						case "7", "07":
 							if !strings.ContainsRune(a.attributes, 'r') {
 								a.attributes += "r"
 							}
 						case "9", "09":
 							if !strings.ContainsRune(a.attributes, 's') {
 								a.attributes += "s"
 							}
 						case "22":
 							if i := strings.IndexRune(a.attributes, 'b'); i >= 0 {
 								a.attributes = a.attributes[:i] + a.attributes[i+1:]
 							}
 							if i := strings.IndexRune(a.attributes, 'd'); i >= 0 {
 								a.attributes = a.attributes[:i] + a.attributes[i+1:]
 							}
 						case "23":
 							if i := strings.IndexRune(a.attributes, 'i'); i >= 0 {
 								a.attributes = a.attributes[:i] + a.attributes[i+1:]
 							}
 						case "24":
 							if i := strings.IndexRune(a.attributes, 'u'); i >= 0 {
 								a.attributes = a.attributes[:i] + a.attributes[i+1:]
 							}
 						case "25":
 							if i := strings.IndexRune(a.attributes, 'l'); i >= 0 {
 								a.attributes = a.attributes[:i] + a.attributes[i+1:]
 							}
 						case "27":
 							if i := strings.IndexRune(a.attributes, 'r'); i >= 0 {
 								a.attributes = a.attributes[:i] + a.attributes[i+1:]
 							}
 						case "29":
 							if i := strings.IndexRune(a.attributes, 's'); i >= 0 {
 								a.attributes = a.attributes[:i] + a.attributes[i+1:]
 							}
 						case "30", "31", "32", "33", "34", "35", "36", "37":
 							colorNumber, _ := strconv.Atoi(field)
 							foreground = lookupColor(colorNumber - 30)
 						case "39":
 							foreground = "-"
 						case "40", "41", "42", "43", "44", "45", "46", "47":
 							colorNumber, _ := strconv.Atoi(field)
 							background = lookupColor(colorNumber - 40)
 						case "49":
 							background = "-"
 						case "90", "91", "92", "93", "94", "95", "96", "97":
 							colorNumber, _ := strconv.Atoi(field)
 							foreground = lookupColor(colorNumber - 82)
 						case "100", "101", "102", "103", "104", "105", "106", "107":
 							colorNumber, _ := strconv.Atoi(field)
 							background = lookupColor(colorNumber - 92)
 						case "38", "48":
 							var color string
 							if len(fields) > index+1 {
 								if fields[index+1] == "5" && len(fields) > index+2 { // 8-bit colors.
 									colorNumber, _ := strconv.Atoi(fields[index+2])
 									if colorNumber <= 15 {
 										color = lookupColor(colorNumber)
 									} else if colorNumber <= 231 {
 										red := (colorNumber - 16) / 36
 										green := ((colorNumber - 16) / 6) % 6
 										blue := (colorNumber - 16) % 6
 										color = fmt.Sprintf("#%02x%02x%02x", 255*red/5, 255*green/5, 255*blue/5)
 									} else if colorNumber <= 255 {
 										grey := 255 * (colorNumber - 232) / 23
 										color = fmt.Sprintf("#%02x%02x%02x", grey, grey, grey)
 									}
 								} else if fields[index+1] == "2" && len(fields) > index+4 { // 24-bit colors.
 									red, _ := strconv.Atoi(fields[index+2])
 									green, _ := strconv.Atoi(fields[index+3])
 									blue, _ := strconv.Atoi(fields[index+4])
 									color = fmt.Sprintf("#%02x%02x%02x", red, green, blue)
 								}
 							}
 							if len(color) > 0 {
 								if field == "38" {
 									foreground = color
 								} else {
 									background = color
 								}
 							}
 							break FieldLoop
 						}
 					}
 					var colon string
 					if len(a.attributes) > 1 && a.attributes[0] == '-' {
 						a.attributes = a.attributes[1:]
 					}
 					if len(a.attributes) > 0 {
 						colon = ":"
 					}
 					if len(foreground) > 0 || len(background) > 0 || len(a.attributes) > 0 {
 						fmt.Fprintf(a.buffer, "[%s:%s%s%s]", foreground, background, colon, a.attributes)
 					}
 				}
 				a.state = ansiText
 			default: // Undefined byte.
 				a.state = ansiText // Abort CSI.
 			}
 
 			// We just entered a substring/command sequence.
 		case ansiSubstring:
 			if r == 27 { // Most likely the end of the substring.
 				a.state = ansiEscape
 			} // Ignore all other characters.
 
 			// "ansiText" and all others.
 		default:
 			if r == 27 {
 				// This is the start of an escape sequence.
 				a.state = ansiEscape
 			} else {
 				// Just a regular rune. Send to buffer.
 				if _, err := a.buffer.WriteRune(r); err != nil {
 					return 0, err
 				}
 			}
 		}
 	}
 
 	// Write buffer to target writer.
 	n, err := a.buffer.WriteTo(a.Writer)
 	if err != nil {
 		return int(n), err
 	}
 	return len(text), nil
 }
 
 // TranslateANSI replaces ANSI escape sequences found in the provided string
 // with tview's style tags and returns the resulting string.
 func TranslateANSI(text string) string {
 	var buffer bytes.Buffer
 	writer := ANSIWriter(&buffer)
 	writer.Write([]byte(text))
 	return buffer.String()
 }