Writer returns number of input bytes consumed, not output bytes written.

writer.go

@@ -67,7 +67,7 @@ func (w *Writer) Colorize() {
 }
 
 // Write writes in to w, prefixing it with an SGR escape sequence, if
-// necessary.  It returns the total number of bytes written and any
+// necessary.  It returns the number of input bytes consumed and any
 // error encountered while writing.
 func (w *Writer) Write(in []byte) (int, error) {
 	p := w.s // primary Writer
@@ -125,7 +125,12 @@ func (w *Writer) Write(in []byte) (int, error) {
 		modes = append(append([]byte{'\033', '['}, modes...), 'm')
 		in = append(modes, in...)
 	}
-	return p.w.Write(in)
+	n, err := p.w.Write(in)
+	n -= len(modes)
+	if n < 0 {
+		n = 0
+	}
+	return n, err
 }
 
 // WriteString writes in to w, prefixing it with an SGR escape sequence, if
@@ -142,22 +147,23 @@ func (w *Writer) WriteString(in string) (int, error) {
 	if err != nil {
 		return n, err
 	}
-	n2, err := old.Set()
-	return n + n2, err
-
+	return n, old.Set()
 }
 
 // Set writes, if necessary, the SGR escape sequence to w to set the current
 // graphics mode.
-func (w *Writer) Set() (int, error) { return w.Write(nil) }
+func (w *Writer) Set() error {
+	_, err := w.Write(nil)
+	return err
+}
 
 // ForceSet writes the SGR escape sequence to w
-func (w *Writer) ForceSet() (int, error) {
+func (w *Writer) ForceSet() error {
 	w = w.copy()
 	w.s.bg = -1
 	w.s.fg = -1
 	w.s.intensity = -1
-	return w.Write(nil)
+	return w.Set()
 }
 
 // SetColor returns a Writer that sets the drawing color to color.

writer_test.go

@@ -6,6 +6,7 @@ package ansi
 
 import (
 	"bytes"
+	"io"
 	"testing"
 )
 
@@ -168,3 +169,29 @@ func TestPrint(t *testing.T) {
 		t.Errorf("Println got %q, want %q", got, want)
 	}
 }
+
+// Writer that "writes" (discards) a limited number of bytes.
+type LimitWriter struct{ limit int }
+
+func (w *LimitWriter) Write(p []byte) (n int, err error) {
+	if w.limit < len(p) {
+		n = w.limit
+		w.limit = 0
+		err = io.EOF
+		return
+	}
+	w.limit -= len(p)
+	return len(p), nil
+}
+
+func TestLength(t *testing.T) {
+	var buf bytes.Buffer
+	if n, err := NewWriter(&buf).Red().Print("hello"); err != nil || n != len("hello") {
+		t.Errorf("Write returned (%v, %v), want (%v, %v)", n, err, len("hello"), nil)
+	}
+
+	limited := LimitWriter{8} // <esc>[31mhello gets truncated after the first l
+	if n, err := NewWriter(&limited).Red().Print("hello"); err != io.EOF || n != len("hel") {
+		t.Errorf("Write returned (%v, %v), want (%v, %v)", n, err, len("hel"), io.EOF)
+	}
+}