fix: include escape codes in byte counts from Print, Printf

Color.Print and Color.Printf return the number of bytes written, but only
counted the bytes of the formatted operands. The SGR escape sequence written
before them and the reset sequence written after were silently dropped from
the count, so the returned n was always short whenever color output was
active.

Fprint and Fprintf had the identical bug and were already fixed in cbda2c3d
by switching from the void Set/unset() pair to setWriter/unsetWriter, which
return the byte counts of the escape sequences so they can be accumulated.
That fix was never applied to Print and Printf, which still call Set() and
unset() and discard what they write.

Give Print and Printf the same setWriter/unsetWriter treatment as Fprint and
Fprintf. The private unset() method has no remaining callers after this
change and is removed.

Add TestColor_Print and TestColor_Printf, mirroring the existing
TestColor_Fprint/TestColor_Fprintf byte-count assertions.
This commit is contained in:
dualfroz committed 2026-09-04 08:25:41 +02:00
1 parent baec7b1239
commit c92652a003
2 files changed
+71 -14

No files matched your search

+26 -14
View File
@@ -235,14 +235,6 @@ func (c *Color) Set() *Color {
return c
}
func (c *Color) unset() {
if c.isNoColorSet() {
return
}
Unset()
}
// SetWriter is used to set the SGR sequence with the given io.Writer. This is
// a low-level function, and users should use the higher-level functions, such
// as color.Fprint, color.Print, etc.
@@ -308,10 +300,20 @@ func (c *Color) Fprint(w io.Writer, a ...interface{}) (n int, err error) {
// encountered. This is the standard fmt.Print() method wrapped with the given
// color.
func (c *Color) Print(a ...interface{}) (n int, err error) {
c.Set()
defer c.unset()
n, err = c.setWriter(Output)
if err != nil {
return n, err
}
return fmt.Fprint(Output, a...)
nn, err := fmt.Fprint(Output, a...)
n += nn
if err != nil {
return
}
nn, err = c.unsetWriter(Output)
n += nn
return n, err
}
// Fprintf formats according to a format specifier and writes to w.
@@ -339,10 +341,20 @@ func (c *Color) Fprintf(w io.Writer, format string, a ...interface{}) (n int, er
// It returns the number of bytes written and any write error encountered.
// This is the standard fmt.Printf() method wrapped with the given color.
func (c *Color) Printf(format string, a ...interface{}) (n int, err error) {
c.Set()
defer c.unset()
n, err = c.setWriter(Output)
if err != nil {
return n, err
}
return fmt.Fprintf(Output, format, a...)
nn, err := fmt.Fprintf(Output, format, a...)
n += nn
if err != nil {
return
}
nn, err = c.unsetWriter(Output)
n += nn
return n, err
}
// Fprintln formats using the default formats for its operands and writes to w.
+45
View File
@@ -555,6 +555,51 @@ func TestColor_Fprintf(t *testing.T) {
}
}
func TestColor_Print(t *testing.T) {
rb := new(strings.Builder)
oldOutput, oldNoColor := Output, NoColor
Output, NoColor = rb, false
t.Cleanup(func() { Output, NoColor = oldOutput, oldNoColor })
c := New(FgRed)
n, err := c.Print("foo", "bar")
if err != nil {
t.Errorf("Print error: %v", err)
}
got := rb.String()
want := "\x1b[31mfoobar\x1b[0m"
if want != got {
t.Errorf("Print error\n\nwant: %q\n got: %q", want, got)
}
if n != len(got) {
t.Errorf("Print byte count does not match actual bytes written\n\nwant: %d\n got: %d", len(got), n)
}
}
func TestColor_Printf(t *testing.T) {
rb := new(strings.Builder)
oldOutput, oldNoColor := Output, NoColor
Output, NoColor = rb, false
t.Cleanup(func() { Output, NoColor = oldOutput, oldNoColor })
c := New(FgRed)
n, err := c.Printf("%-7s %-7s %5d\n", "hello", "world", 123)
if err != nil {
t.Errorf("Printf error: %v", err)
}
want := "\x1b[31mhello world 123\n\x1b[0m"
got := rb.String()
if want != got {
t.Errorf("Printf error\n\nwant: %q\n got: %q", want, got)
}
if n != len(got) {
t.Errorf("Printf byte count does not match actual bytes written\n\nwant: %d\n got: %d", len(got), n)
}
}
func TestColor_Fprintln_Newline(t *testing.T) {
rb := new(bytes.Buffer)
c := New(FgRed)