output_windows.go raw

   1  package ansi
   2  
   3  import (
   4  	"bytes"
   5  	"fmt"
   6  	"io"
   7  	"os"
   8  	"strconv"
   9  	"strings"
  10  	"syscall"
  11  	"unsafe"
  12  
  13  	"github.com/mattn/go-isatty"
  14  )
  15  
  16  var (
  17  	singleArgFunctions = map[rune]func(int){
  18  		'A': CursorUp,
  19  		'B': CursorDown,
  20  		'C': CursorForward,
  21  		'D': CursorBack,
  22  		'E': CursorNextLine,
  23  		'F': CursorPreviousLine,
  24  		'G': CursorHorizontalAbsolute,
  25  	}
  26  )
  27  
  28  const (
  29  	foregroundBlue      = 0x1
  30  	foregroundGreen     = 0x2
  31  	foregroundRed       = 0x4
  32  	foregroundIntensity = 0x8
  33  	foregroundMask      = (foregroundRed | foregroundBlue | foregroundGreen | foregroundIntensity)
  34  	backgroundBlue      = 0x10
  35  	backgroundGreen     = 0x20
  36  	backgroundRed       = 0x40
  37  	backgroundIntensity = 0x80
  38  	backgroundMask      = (backgroundRed | backgroundBlue | backgroundGreen | backgroundIntensity)
  39  )
  40  
  41  type Writer struct {
  42  	out     io.Writer
  43  	handle  syscall.Handle
  44  	orgAttr word
  45  }
  46  
  47  func NewAnsiStdout() io.Writer {
  48  	var csbi consoleScreenBufferInfo
  49  	out := os.Stdout
  50  	if !isatty.IsTerminal(out.Fd()) {
  51  		return out
  52  	}
  53  	handle := syscall.Handle(out.Fd())
  54  	procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
  55  	return &Writer{out: out, handle: handle, orgAttr: csbi.attributes}
  56  }
  57  
  58  func NewAnsiStderr() io.Writer {
  59  	var csbi consoleScreenBufferInfo
  60  	out := os.Stderr
  61  	if !isatty.IsTerminal(out.Fd()) {
  62  		return out
  63  	}
  64  	handle := syscall.Handle(out.Fd())
  65  	procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
  66  	return &Writer{out: out, handle: handle, orgAttr: csbi.attributes}
  67  }
  68  
  69  func (w *Writer) Write(data []byte) (n int, err error) {
  70  	r := bytes.NewReader(data)
  71  
  72  	for {
  73  		ch, size, err := r.ReadRune()
  74  		if err != nil {
  75  			break
  76  		}
  77  		n += size
  78  
  79  		switch ch {
  80  		case '\x1b':
  81  			size, err = w.handleEscape(r)
  82  			n += size
  83  			if err != nil {
  84  				break
  85  			}
  86  		default:
  87  			fmt.Fprint(w.out, string(ch))
  88  		}
  89  	}
  90  	return
  91  }
  92  
  93  func (w *Writer) handleEscape(r *bytes.Reader) (n int, err error) {
  94  	buf := make([]byte, 0, 10)
  95  	buf = append(buf, "\x1b"...)
  96  
  97  	// Check '[' continues after \x1b
  98  	ch, size, err := r.ReadRune()
  99  	if err != nil {
 100  		fmt.Fprint(w.out, string(buf))
 101  		return
 102  	}
 103  	n += size
 104  	if ch != '[' {
 105  		fmt.Fprint(w.out, string(buf))
 106  		return
 107  	}
 108  
 109  	// Parse escape code
 110  	var code rune
 111  	argBuf := make([]byte, 0, 10)
 112  	for {
 113  		ch, size, err = r.ReadRune()
 114  		if err != nil {
 115  			fmt.Fprint(w.out, string(buf))
 116  			return
 117  		}
 118  		n += size
 119  		if ('a' <= ch && ch <= 'z') || ('A' <= ch && ch <= 'Z') {
 120  			code = ch
 121  			break
 122  		}
 123  		argBuf = append(argBuf, string(ch)...)
 124  	}
 125  
 126  	w.applyEscapeCode(buf, string(argBuf), code)
 127  	return
 128  }
 129  
 130  func (w *Writer) applyEscapeCode(buf []byte, arg string, code rune) {
 131  	switch arg + string(code) {
 132  	case "?25h":
 133  		CursorShow()
 134  		return
 135  	case "?25l":
 136  		CursorHide()
 137  		return
 138  	}
 139  
 140  	if f, ok := singleArgFunctions[code]; ok {
 141  		if n, err := strconv.Atoi(arg); err == nil {
 142  			f(n)
 143  			return
 144  		}
 145  	}
 146  
 147  	switch code {
 148  	case 'm':
 149  		w.applySelectGraphicRendition(arg)
 150  	default:
 151  		buf = append(buf, string(code)...)
 152  		fmt.Fprint(w.out, string(buf))
 153  	}
 154  }
 155  
 156  // Original implementation: https://github.com/mattn/go-colorable
 157  // FIXME: Fallback to original go-colorable to avoid duplicate implementation
 158  func (w *Writer) applySelectGraphicRendition(arg string) {
 159  	if arg == "" {
 160  		procSetConsoleTextAttribute.Call(uintptr(w.handle), uintptr(w.orgAttr))
 161  		return
 162  	}
 163  
 164  	var csbi consoleScreenBufferInfo
 165  	procGetConsoleScreenBufferInfo.Call(uintptr(w.handle), uintptr(unsafe.Pointer(&csbi)))
 166  	attr := csbi.attributes
 167  
 168  	for _, param := range strings.Split(arg, ";") {
 169  		n, err := strconv.Atoi(param)
 170  		if err != nil {
 171  			continue
 172  		}
 173  
 174  		switch {
 175  		case n == 0 || n == 100:
 176  			attr = w.orgAttr
 177  		case 1 <= n && n <= 5:
 178  			attr |= foregroundIntensity
 179  		case 30 <= n && n <= 37:
 180  			attr = (attr & backgroundMask)
 181  			if (n-30)&1 != 0 {
 182  				attr |= foregroundRed
 183  			}
 184  			if (n-30)&2 != 0 {
 185  				attr |= foregroundGreen
 186  			}
 187  			if (n-30)&4 != 0 {
 188  				attr |= foregroundBlue
 189  			}
 190  		case 40 <= n && n <= 47:
 191  			attr = (attr & foregroundMask)
 192  			if (n-40)&1 != 0 {
 193  				attr |= backgroundRed
 194  			}
 195  			if (n-40)&2 != 0 {
 196  				attr |= backgroundGreen
 197  			}
 198  			if (n-40)&4 != 0 {
 199  				attr |= backgroundBlue
 200  			}
 201  		}
 202  	}
 203  
 204  	procSetConsoleTextAttribute.Call(uintptr(w.handle), uintptr(attr))
 205  }
 206