colorwriter_windows.go

Documentation: github.com/go-kit/log/term

     1  // The code in this file is adapted from github.com/mattn/go-colorable.
     2  
     3  package term
     4  
     5  import (
     6  	"bytes"
     7  	"fmt"
     8  	"io"
     9  	"strconv"
    10  	"strings"
    11  	"syscall"
    12  	"unsafe"
    13  )
    14  
    15  type colorWriter struct {
    16  	out     io.Writer
    17  	handle  syscall.Handle
    18  	lastbuf bytes.Buffer
    19  	oldattr word
    20  }
    21  
    22  // NewColorWriter returns an io.Writer that writes to w and provides cross
    23  // platform support for ANSI color codes. If w is not a terminal it is
    24  // returned unmodified.
    25  func NewColorWriter(w io.Writer) io.Writer {
    26  	if !IsConsole(w) {
    27  		return w
    28  	}
    29  
    30  	var csbi consoleScreenBufferInfo
    31  	handle := syscall.Handle(w.(fder).Fd())
    32  	procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
    33  
    34  	return &colorWriter{
    35  		out:     w,
    36  		handle:  handle,
    37  		oldattr: csbi.attributes,
    38  	}
    39  }
    40  
    41  func (w *colorWriter) Write(data []byte) (n int, err error) {
    42  	var csbi consoleScreenBufferInfo
    43  	procGetConsoleScreenBufferInfo.Call(uintptr(w.handle), uintptr(unsafe.Pointer(&csbi)))
    44  
    45  	er := bytes.NewBuffer(data)
    46  loop:
    47  	for {
    48  		r1, _, _ := procGetConsoleScreenBufferInfo.Call(uintptr(w.handle), uintptr(unsafe.Pointer(&csbi)))
    49  		if r1 == 0 {
    50  			break loop
    51  		}
    52  
    53  		c1, _, err := er.ReadRune()
    54  		if err != nil {
    55  			break loop
    56  		}
    57  		if c1 != 0x1b {
    58  			fmt.Fprint(w.out, string(c1))
    59  			continue
    60  		}
    61  		c2, _, err := er.ReadRune()
    62  		if err != nil {
    63  			w.lastbuf.WriteRune(c1)
    64  			break loop
    65  		}
    66  		if c2 != 0x5b {
    67  			w.lastbuf.WriteRune(c1)
    68  			w.lastbuf.WriteRune(c2)
    69  			continue
    70  		}
    71  
    72  		var buf bytes.Buffer
    73  		var m rune
    74  		for {
    75  			c, _, err := er.ReadRune()
    76  			if err != nil {
    77  				w.lastbuf.WriteRune(c1)
    78  				w.lastbuf.WriteRune(c2)
    79  				w.lastbuf.Write(buf.Bytes())
    80  				break loop
    81  			}
    82  			if ('a' <= c && c <= 'z') || ('A' <= c && c <= 'Z') || c == '@' {
    83  				m = c
    84  				break
    85  			}
    86  			buf.Write([]byte(string(c)))
    87  		}
    88  
    89  		switch m {
    90  		case 'm':
    91  			attr := csbi.attributes
    92  			cs := buf.String()
    93  			if cs == "" {
    94  				procSetConsoleTextAttribute.Call(uintptr(w.handle), uintptr(w.oldattr))
    95  				continue
    96  			}
    97  			token := strings.Split(cs, ";")
    98  			intensityMode := word(0)
    99  			for _, ns := range token {
   100  				if n, err = strconv.Atoi(ns); err == nil {
   101  					switch {
   102  					case n == 0:
   103  						attr = w.oldattr
   104  					case n == 1:
   105  						attr |= intensityMode
   106  					case 30 <= n && n <= 37:
   107  						attr = (attr & backgroundMask)
   108  						if (n-30)&1 != 0 {
   109  							attr |= foregroundRed
   110  						}
   111  						if (n-30)&2 != 0 {
   112  							attr |= foregroundGreen
   113  						}
   114  						if (n-30)&4 != 0 {
   115  							attr |= foregroundBlue
   116  						}
   117  						intensityMode = foregroundIntensity
   118  					case n == 39: // reset foreground color
   119  						attr &= backgroundMask
   120  						attr |= w.oldattr & foregroundMask
   121  					case 40 <= n && n <= 47:
   122  						attr = (attr & foregroundMask)
   123  						if (n-40)&1 != 0 {
   124  							attr |= backgroundRed
   125  						}
   126  						if (n-40)&2 != 0 {
   127  							attr |= backgroundGreen
   128  						}
   129  						if (n-40)&4 != 0 {
   130  							attr |= backgroundBlue
   131  						}
   132  						intensityMode = backgroundIntensity
   133  					case n == 49: // reset background color
   134  						attr &= foregroundMask
   135  						attr |= w.oldattr & backgroundMask
   136  					}
   137  					procSetConsoleTextAttribute.Call(uintptr(w.handle), uintptr(attr))
   138  				}
   139  			}
   140  		}
   141  	}
   142  	return len(data) - w.lastbuf.Len(), nil
   143  }
   144  
   145  var (
   146  	procGetConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo")
   147  	procSetConsoleTextAttribute    = kernel32.NewProc("SetConsoleTextAttribute")
   148  )
   149  
   150  const (
   151  	foregroundBlue      = 0x1
   152  	foregroundGreen     = 0x2
   153  	foregroundRed       = 0x4
   154  	foregroundIntensity = 0x8
   155  	foregroundMask      = (foregroundRed | foregroundBlue | foregroundGreen | foregroundIntensity)
   156  	backgroundBlue      = 0x10
   157  	backgroundGreen     = 0x20
   158  	backgroundRed       = 0x40
   159  	backgroundIntensity = 0x80
   160  	backgroundMask      = (backgroundRed | backgroundBlue | backgroundGreen | backgroundIntensity)
   161  )
   162  
   163  type (
   164  	wchar uint16 //lint:ignore U1000 unused
   165  	short int16  //lint:ignore U1000 unused
   166  	dword uint32 //lint:ignore U1000 unused
   167  	word  uint16 //lint:ignore U1000 unused
   168  )
   169  
   170  type coord struct {
   171  	x short //lint:ignore U1000 unused
   172  	y short //lint:ignore U1000 unused
   173  }
   174  
   175  type smallRect struct {
   176  	left   short //lint:ignore U1000 unused
   177  	top    short //lint:ignore U1000 unused
   178  	right  short //lint:ignore U1000 unused
   179  	bottom short //lint:ignore U1000 unused
   180  }
   181  
   182  type consoleScreenBufferInfo struct {
   183  	size              coord
   184  	cursorPosition    coord
   185  	attributes        word
   186  	window            smallRect
   187  	maximumWindowSize coord
   188  }
   189
View as plain text