From 886f222e9816bbffb7304d27a45ba91d7586861b Mon Sep 17 00:00:00 2001
From: Matthieu Dumont <5095856+Jerska@users.noreply.github.com>
Date: Wed, 16 Feb 2022 18:44:13 +0100
Subject: [PATCH] Encoder: add SetAddExtraNewline option to json.Encoder (#119)

* Encoder: add SetAddExtraNewline option to json.Encoder

* feat: rename AddExtraNewline to AppendNewline

* refactor: change AppendNewline to be part of AppendFlags

* feat: make appendNewline private
---
 json/json.go      | 23 +++++++++++++++++++++--
 json/json_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/json/json.go b/json/json.go
index 9d259cc..47f3ba1 100644
--- a/json/json.go
+++ b/json/json.go
@@ -70,6 +70,11 @@ const (
 	// checking of raw messages. It should only be used if the values are
 	// known to be valid json (e.g., they were created by json.Unmarshal).
 	TrustRawMessage
+
+	// appendNewline is a formatting flag to enable the addition of a newline
+	// in Encode (this matches the behavior of the standard encoding/json
+	// package).
+	appendNewline
 )
 
 // ParseFlags is a type used to represent configuration options that can be
@@ -480,7 +485,9 @@ type Encoder struct {
 }
 
 // NewEncoder is documented at https://golang.org/pkg/encoding/json/#NewEncoder
-func NewEncoder(w io.Writer) *Encoder { return &Encoder{writer: w, flags: EscapeHTML | SortMapKeys} }
+func NewEncoder(w io.Writer) *Encoder {
+	return &Encoder{writer: w, flags: EscapeHTML | SortMapKeys | appendNewline}
+}
 
 // Encode is documented at https://golang.org/pkg/encoding/json/#Encoder.Encode
 func (enc *Encoder) Encode(v interface{}) error {
@@ -498,7 +505,9 @@ func (enc *Encoder) Encode(v interface{}) error {
 		return err
 	}
 
-	buf.data = append(buf.data, '\n')
+	if (enc.flags & appendNewline) != 0 {
+		buf.data = append(buf.data, '\n')
+	}
 	b := buf.data
 
 	if enc.prefix != "" || enc.indent != "" {
@@ -556,6 +565,16 @@ func (enc *Encoder) SetTrustRawMessage(on bool) {
 	}
 }
 
+// SetAppendNewline is an extension to the standard encoding/json package which
+// allows the program to toggle the addition of a newline in Encode on or off.
+func (enc *Encoder) SetAppendNewline(on bool) {
+	if on {
+		enc.flags |= appendNewline
+	} else {
+		enc.flags &= ^appendNewline
+	}
+}
+
 var encoderBufferPool = sync.Pool{
 	New: func() interface{} { return &encoderBuffer{data: make([]byte, 0, 4096)} },
 }
diff --git a/json/json_test.go b/json/json_test.go
index 63ec861..6840b85 100644
--- a/json/json_test.go
+++ b/json/json_test.go
@@ -1669,6 +1669,44 @@ func TestSetTrustRawMessage(t *testing.T) {
 	}
 }
 
+func TestSetAppendNewline(t *testing.T) {
+	buf := &bytes.Buffer{}
+	enc := NewEncoder(buf)
+
+	m := "value"
+
+	// Default encoding adds an extra newline
+	if err := enc.Encode(m); err != nil {
+		t.Error(err)
+	}
+	b := buf.Bytes()
+	exp := []byte(`"value"`)
+	exp = append(exp, '\n')
+	if bytes.Compare(exp, b) != 0 {
+		t.Error(
+			"unexpected encoding:",
+			"expected", exp,
+			"got", b,
+		)
+	}
+
+	// With SetAppendNewline(false), there shouldn't be a newline in the output
+	buf.Reset()
+	enc.SetAppendNewline(false)
+	if err := enc.Encode(m); err != nil {
+		t.Error(err)
+	}
+	b = buf.Bytes()
+	exp = []byte(`"value"`)
+	if bytes.Compare(exp, b) != 0 {
+		t.Error(
+			"unexpected encoding:",
+			"expected", exp,
+			"got", b,
+		)
+	}
+}
+
 func TestEscapeString(t *testing.T) {
 	b := Escape(`value`)
 	x := []byte(`"value"`)