output_windows.go 4.5 KB

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