From aa57fa9006134bc0b23eb4e63c44ff5684af7f48 Mon Sep 17 00:00:00 2001
From: kiyon <kiyonlin@163.com>
Date: Wed, 21 Oct 2020 14:47:29 +0800
Subject: [PATCH] Add a new method `Flush` to reuse writer (#18)

---
 padding/padding.go      | 37 ++++++++++++++++---------
 padding/padding_test.go | 60 ++++++++++++++++++++++++++++++++++++++---
 2 files changed, 81 insertions(+), 16 deletions(-)

diff --git a/padding/padding.go b/padding/padding.go
index b21ba05..707ec46 100644
--- a/padding/padding.go
+++ b/padding/padding.go
@@ -6,7 +6,6 @@ import (
 	"strings"
 
 	"github.com/mattn/go-runewidth"
-
 	"github.com/muesli/reflow/ansi"
 )
 
@@ -18,6 +17,7 @@ type Writer struct {
 
 	ansiWriter *ansi.Writer
 	buf        bytes.Buffer
+	cache      bytes.Buffer
 	lineLen    int
 	ansi       bool
 }
@@ -48,7 +48,7 @@ func NewWriterPipe(forward io.Writer, width uint, paddingFunc PaddingFunc) *Writ
 func Bytes(b []byte, width uint) []byte {
 	f := NewWriter(width, nil)
 	_, _ = f.Write(b)
-	f.Close()
+	_ = f.Flush()
 
 	return f.Bytes()
 }
@@ -110,23 +110,34 @@ func (w *Writer) pad() error {
 	return nil
 }
 
-// Close will finish the padding operation. Always call it before trying to
-// retrieve the final result.
-func (w *Writer) Close() error {
-	// don't pad empty trailing lines
-	if w.lineLen == 0 {
-		return nil
-	}
-
-	return w.pad()
+// Close will finish the padding operation.
+func (w *Writer) Close() (err error) {
+	return w.Flush()
 }
 
 // Bytes returns the padded result as a byte slice.
 func (w *Writer) Bytes() []byte {
-	return w.buf.Bytes()
+	return w.cache.Bytes()
 }
 
 // String returns the padded result as a string.
 func (w *Writer) String() string {
-	return w.buf.String()
+	return w.cache.String()
+}
+
+// Flush will finish the padding operation. Always call it before trying to
+// retrieve the final result.
+func (w *Writer) Flush() (err error) {
+	if w.lineLen != 0 {
+		if err = w.pad(); err != nil {
+			return
+		}
+	}
+
+	w.cache.Reset()
+	_, err = w.buf.WriteTo(&w.cache)
+	w.lineLen = 0
+	w.ansi = false
+
+	return
 }
diff --git a/padding/padding_test.go b/padding/padding_test.go
index bc7d942..5327100 100644
--- a/padding/padding_test.go
+++ b/padding/padding_test.go
@@ -56,7 +56,10 @@ func TestPadding(t *testing.T) {
 		if err != nil {
 			t.Error(err)
 		}
-		f.Close()
+
+		if err := f.Close(); err != nil {
+			t.Error(err)
+		}
 
 		if f.String() != tc.Expected {
 			t.Errorf("Test %d, expected:\n\n`%s`\n\nActual Output:\n\n`%s`", i, tc.Expected, f.String())
@@ -77,7 +80,9 @@ func TestPaddingWriter(t *testing.T) {
 	if err != nil {
 		t.Error(err)
 	}
-	f.Close()
+	if err := f.Close(); err != nil {
+		t.Error(err)
+	}
 
 	exp := "foo   \nbar   "
 	if f.String() != exp {
@@ -114,7 +119,9 @@ func TestNewWriterPipe(t *testing.T) {
 	if _, err := f.Write([]byte("foobar")); err != nil {
 		t.Error(err)
 	}
-	f.Close()
+	if err := f.Close(); err != nil {
+		t.Error(err)
+	}
 
 	actual := b.String()
 	expected := "foobar    "
@@ -142,6 +149,53 @@ func TestWriter_pad(t *testing.T) {
 	}
 }
 
+func TestWriter_Flush(t *testing.T) {
+	t.Parallel()
+
+	f := NewWriter(6, nil)
+
+	_, err := f.Write([]byte("foo"))
+	if err != nil {
+		t.Error(err)
+	}
+
+	if err := f.Flush(); err != nil {
+		t.Error(err)
+	}
+
+	exp := "foo   "
+	if f.String() != exp {
+		t.Errorf("expected:\n\n`%s`\n\nActual Output:\n\n`%s`", exp, f.String())
+	}
+
+	_, err = f.Write([]byte("bar"))
+	if err != nil {
+		t.Error(err)
+	}
+	if err := f.Flush(); err != nil {
+		t.Error(err)
+	}
+
+	exp = "bar   "
+	if f.String() != exp {
+		t.Errorf("expected:\n\n`%s`\n\nActual Output:\n\n`%s`", exp, f.String())
+	}
+}
+
+func TestWriter_Close(t *testing.T) {
+	t.Parallel()
+
+	f := &Writer{
+		Padding:    6,
+		lineLen:    1,
+		ansiWriter: &ansi.Writer{Forward: fakeWriter{}},
+	}
+
+	if err := f.Close(); err != fakeErr {
+		t.Error(err)
+	}
+}
+
 func TestWriter_Error(t *testing.T) {
 	t.Parallel()