ansi.go

Documentation: code.rocketnine.space/tslocum/cview

     1  package cview
     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 cview 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 cview 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.IndexRune(a.attributes, 'b') < 0 {
   132  								a.attributes += "b"
   133  							}
   134  						case "2", "02":
   135  							if strings.IndexRune(a.attributes, 'd') < 0 {
   136  								a.attributes += "d"
   137  							}
   138  						case "3", "03":
   139  							if strings.IndexRune(a.attributes, 'i') < 0 {
   140  								a.attributes += "i"
   141  							}
   142  						case "4", "04":
   143  							if strings.IndexRune(a.attributes, 'u') < 0 {
   144  								a.attributes += "u"
   145  							}
   146  						case "5", "05":
   147  							if strings.IndexRune(a.attributes, 'l') < 0 {
   148  								a.attributes += "l"
   149  							}
   150  						case "7", "07":
   151  							if strings.IndexRune(a.attributes, 'r') < 0 {
   152  								a.attributes += "r"
   153  							}
   154  						case "9", "09":
   155  							if strings.IndexRune(a.attributes, 's') < 0 {
   156  								a.attributes += "s"
   157  							}
   158  						case "22":
   159  							if i := strings.IndexRune(a.attributes, 'b'); i >= 0 {
   160  								a.attributes = a.attributes[:i] + a.attributes[i+1:]
   161  							}
   162  							if i := strings.IndexRune(a.attributes, 'd'); i >= 0 {
   163  								a.attributes = a.attributes[:i] + a.attributes[i+1:]
   164  							}
   165  						case "24":
   166  							if i := strings.IndexRune(a.attributes, 'u'); i >= 0 {
   167  								a.attributes = a.attributes[:i] + a.attributes[i+1:]
   168  							}
   169  						case "25":
   170  							if i := strings.IndexRune(a.attributes, 'l'); i >= 0 {
   171  								a.attributes = a.attributes[:i] + a.attributes[i+1:]
   172  							}
   173  						case "30", "31", "32", "33", "34", "35", "36", "37":
   174  							colorNumber, _ := strconv.Atoi(field)
   175  							foreground = lookupColor(colorNumber - 30)
   176  						case "39":
   177  							foreground = "-"
   178  						case "40", "41", "42", "43", "44", "45", "46", "47":
   179  							colorNumber, _ := strconv.Atoi(field)
   180  							background = lookupColor(colorNumber - 40)
   181  						case "49":
   182  							background = "-"
   183  						case "90", "91", "92", "93", "94", "95", "96", "97":
   184  							colorNumber, _ := strconv.Atoi(field)
   185  							foreground = lookupColor(colorNumber - 82)
   186  						case "100", "101", "102", "103", "104", "105", "106", "107":
   187  							colorNumber, _ := strconv.Atoi(field)
   188  							background = lookupColor(colorNumber - 92)
   189  						case "38", "48":
   190  							var color string
   191  							if len(fields) > index+1 {
   192  								if fields[index+1] == "5" && len(fields) > index+2 { // 8-bit colors.
   193  									colorNumber, _ := strconv.Atoi(fields[index+2])
   194  									if colorNumber <= 15 {
   195  										color = lookupColor(colorNumber)
   196  									} else if colorNumber <= 231 {
   197  										red := (colorNumber - 16) / 36
   198  										green := ((colorNumber - 16) / 6) % 6
   199  										blue := (colorNumber - 16) % 6
   200  										color = fmt.Sprintf("#%02x%02x%02x", 255*red/5, 255*green/5, 255*blue/5)
   201  									} else if colorNumber <= 255 {
   202  										grey := 255 * (colorNumber - 232) / 23
   203  										color = fmt.Sprintf("#%02x%02x%02x", grey, grey, grey)
   204  									}
   205  								} else if fields[index+1] == "2" && len(fields) > index+4 { // 24-bit colors.
   206  									red, _ := strconv.Atoi(fields[index+2])
   207  									green, _ := strconv.Atoi(fields[index+3])
   208  									blue, _ := strconv.Atoi(fields[index+4])
   209  									color = fmt.Sprintf("#%02x%02x%02x", red, green, blue)
   210  								}
   211  							}
   212  							if len(color) > 0 {
   213  								if field == "38" {
   214  									foreground = color
   215  								} else {
   216  									background = color
   217  								}
   218  							}
   219  							break FieldLoop
   220  						}
   221  					}
   222  					var colon string
   223  					if len(a.attributes) > 0 {
   224  						colon = ":"
   225  					}
   226  					if len(foreground) > 0 || len(background) > 0 || len(a.attributes) > 0 {
   227  						fmt.Fprintf(a.buffer, "[%s:%s%s%s]", foreground, background, colon, a.attributes)
   228  					}
   229  				}
   230  				a.state = ansiText
   231  			default: // Undefined byte.
   232  				a.state = ansiText // Abort CSI.
   233  			}
   234  
   235  			// We just entered a substring/command sequence.
   236  		case ansiSubstring:
   237  			if r == 27 { // Most likely the end of the substring.
   238  				a.state = ansiEscape
   239  			} // Ignore all other characters.
   240  
   241  			// "ansiText" and all others.
   242  		default:
   243  			if r == 27 {
   244  				// This is the start of an escape sequence.
   245  				a.state = ansiEscape
   246  			} else {
   247  				// Just a regular rune. Send to buffer.
   248  				if _, err := a.buffer.WriteRune(r); err != nil {
   249  					return 0, err
   250  				}
   251  			}
   252  		}
   253  	}
   254  
   255  	// Write buffer to target writer.
   256  	n, err := a.buffer.WriteTo(a.Writer)
   257  	if err != nil {
   258  		return int(n), err
   259  	}
   260  	return len(text), nil
   261  }
   262  
   263  // TranslateANSI replaces ANSI escape sequences found in the provided string
   264  // with cview's color tags and returns the resulting string.
   265  func TranslateANSI(text string) string {
   266  	var buffer bytes.Buffer
   267  	writer := ANSIWriter(&buffer)
   268  	writer.Write([]byte(text))
   269  	return buffer.String()
   270  }
   271
View as plain text