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)