From c92652a003feea3e388d8dffde96096b548f980e Mon Sep 17 00:00:00 2001
From: DualFroz <me@dualfroz.com>
Date: Fri, 4 Sep 2026 08:25:41 +0200
Subject: [PATCH] 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.
---
 color.go      | 40 ++++++++++++++++++++++++++--------------
 color_test.go | 45 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 14 deletions(-)

diff --git a/color.go b/color.go
index d3906bf..ba5ede2 100644
--- a/color.go
+++ b/color.go
@@ -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.
diff --git a/color_test.go b/color_test.go
index e22af6b..628c8ca 100644
--- a/color_test.go
+++ b/color_test.go
@@ -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)