瀏覽代碼

Merge pull request #13 from mislav/generic-constructor

Add a generic `NewColorable()` constructor
mattn 9 年之前
父節點
當前提交
9cbef7c353
共有 2 個文件被更改,包括 25 次插入16 次删除
  1. 8 0
      colorable_others.go
  2. 17 16
      colorable_windows.go

+ 8 - 0
colorable_others.go

@@ -7,6 +7,14 @@ import (
 	"os"
 )
 
+func NewColorable(file *os.File) io.Writer {
+	if file == nil {
+		panic("nil passed instead of *os.File to NewColorable()")
+	}
+
+	return file
+}
+
 func NewColorableStdout() io.Writer {
 	return os.Stdout
 }

+ 17 - 16
colorable_windows.go

@@ -68,26 +68,27 @@ type Writer struct {
 	oldattr word
 }
 
-func NewColorableStdout() io.Writer {
-	var csbi consoleScreenBufferInfo
-	out := os.Stdout
-	if !isatty.IsTerminal(out.Fd()) {
-		return out
+func NewColorable(file *os.File) io.Writer {
+	if file == nil {
+		panic("nil passed instead of *os.File to NewColorable()")
+	}
+
+	if isatty.IsTerminal(file.Fd()) {
+		var csbi consoleScreenBufferInfo
+		handle := syscall.Handle(file.Fd())
+		procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
+		return &Writer{out: file, handle: handle, oldattr: csbi.attributes}
+	} else {
+		return file
 	}
-	handle := syscall.Handle(out.Fd())
-	procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
-	return &Writer{out: out, handle: handle, oldattr: csbi.attributes}
+}
+
+func NewColorableStdout() io.Writer {
+	return NewColorable(os.Stdout)
 }
 
 func NewColorableStderr() io.Writer {
-	var csbi consoleScreenBufferInfo
-	out := os.Stderr
-	if !isatty.IsTerminal(out.Fd()) {
-		return out
-	}
-	handle := syscall.Handle(out.Fd())
-	procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
-	return &Writer{out: out, handle: handle, oldattr: csbi.attributes}
+	return NewColorable(os.Stderr)
 }
 
 var color256 = map[int]int{