From 10c1dcc69a0082c36f6ff62c6a55d7aec03a9a78 Mon Sep 17 00:00:00 2001
From: Arvind Thirunarayanan <athirunar@confluent.io>
Date: Mon, 29 Jul 2019 10:32:00 -0500
Subject: [PATCH 1/5] Add preserveFormatting option for comments/whitespace

---
 decode_test.go     |  2 +-
 lex.go             | 43 ++++++++++++++++----
 load.go            | 19 ++++++---
 parser.go          | 35 ++++++++++++++---
 properties.go      | 97 +++++++++++++++++++++++++++++++++++++---------
 properties_test.go |  6 +--
 6 files changed, 162 insertions(+), 40 deletions(-)

diff --git a/decode_test.go b/decode_test.go
index 9431b83..387de8c 100644
--- a/decode_test.go
+++ b/decode_test.go
@@ -278,7 +278,7 @@ func TestDecodeMap(t *testing.T) {
 }
 
 func testDecode(t *testing.T, in string, v, out interface{}) {
-	p, err := parse(in)
+	p, err := parse(in, false)
 	if err != nil {
 		t.Fatalf("got %v want nil", err)
 	}
diff --git a/lex.go b/lex.go
index 367166d..7f376c3 100644
--- a/lex.go
+++ b/lex.go
@@ -60,6 +60,7 @@ type stateFn func(*lexer) stateFn
 
 // lexer holds the state of the scanner.
 type lexer struct {
+	preserveFormatting bool  // whether to scan EOLs/whitespace as part of comments
 	input   string    // the string being scanned
 	state   stateFn   // the next lexing function to enter
 	pos     int       // current position in the input
@@ -162,8 +163,9 @@ func (l *lexer) nextItem() item {
 }
 
 // lex creates a new scanner for the input string.
-func lex(input string) *lexer {
+func lex(input string, preserveFormatting bool) *lexer {
 	l := &lexer{
+		preserveFormatting: preserveFormatting,
 		input: input,
 		items: make(chan item),
 		runes: make([]rune, 0, 32),
@@ -189,14 +191,36 @@ func lexBeforeKey(l *lexer) stateFn {
 		return nil
 
 	case isEOL(r):
-		l.ignore()
-		return lexBeforeKey
+		if l.preserveFormatting {
+			// treat as part of the next key's comments block
+			// add the EOL rune r as the comment's prefix
+			l.appendRune(r)
+			l.backup()
+			return lexComment
+		} else {
+			l.ignore()
+			return lexBeforeKey
+		}
 
 	case isComment(r):
+		if l.preserveFormatting {
+			// treat as part of the next key's comments block
+			// use existing prefix r as the comment's prefix
+			l.appendRune(r)
+		} else {
+			// use "#" as the comment's prefix
+			l.appendRune('#')
+		}
 		return lexComment
 
 	case isWhitespace(r):
-		l.ignore()
+		if l.preserveFormatting {
+			// treat as part of the next key's comments block
+			// add the whitespace rune r as the comment's prefix
+			l.appendRune(r)
+		} else {
+			l.ignore()
+		}
 		return lexBeforeKey
 
 	default:
@@ -207,12 +231,17 @@ func lexBeforeKey(l *lexer) stateFn {
 
 // lexComment scans a comment line. The comment character has already been scanned.
 func lexComment(l *lexer) stateFn {
-	l.acceptRun(whitespace)
-	l.ignore()
+	if ! l.preserveFormatting {
+		l.acceptRun(whitespace)
+		l.ignore()
+	}
 	for {
+
 		switch r := l.next(); {
 		case isEOF(r):
-			l.ignore()
+			if ! l.preserveFormatting {
+				l.ignore()
+			}
 			l.emit(itemEOF)
 			return nil
 		case isEOL(r):
diff --git a/load.go b/load.go
index ab95325..11250de 100644
--- a/load.go
+++ b/load.go
@@ -45,11 +45,16 @@ type Loader struct {
 	// 404 are reported as errors. When set to true, missing files and 404
 	// status codes are not reported as errors.
 	IgnoreMissing bool
+
+	// PreserveFormatting causes the loader to scan whitespace as part of the
+	// comments for the next key. These can then be written out by passing
+	// a similar preserveFormatting argument to WriteCommentWithFormatting.
+	PreserveFormatting bool
 }
 
 // Load reads a buffer into a Properties struct.
 func (l *Loader) LoadBytes(buf []byte) (*Properties, error) {
-	return l.loadBytes(buf, l.Encoding)
+	return l.loadBytes(buf, l.Encoding, l.PreserveFormatting)
 }
 
 // LoadAll reads the content of multiple URLs or files in the given order into
@@ -99,7 +104,7 @@ func (l *Loader) LoadFile(filename string) (*Properties, error) {
 		}
 		return nil, err
 	}
-	return l.loadBytes(data, l.Encoding)
+	return l.loadBytes(data, l.Encoding, l.PreserveFormatting)
 }
 
 // LoadURL reads the content of the URL into a Properties struct.
@@ -142,11 +147,11 @@ func (l *Loader) LoadURL(url string) (*Properties, error) {
 		return nil, fmt.Errorf("properties: invalid content type %s", ct)
 	}
 
-	return l.loadBytes(body, enc)
+	return l.loadBytes(body, enc, l.PreserveFormatting)
 }
 
-func (l *Loader) loadBytes(buf []byte, enc Encoding) (*Properties, error) {
-	p, err := parse(convert(buf, enc))
+func (l *Loader) loadBytes(buf []byte, enc Encoding, preserveFormatting bool) (*Properties, error) {
+	p, err := parse(convert(buf, enc), preserveFormatting)
 	if err != nil {
 		return nil, err
 	}
@@ -178,6 +183,10 @@ func LoadMap(m map[string]string) *Properties {
 	return p
 }
 
+func GetLoader() (*Loader, error) {
+	return &Loader{}, nil
+}
+
 // LoadFile reads a file into a Properties struct.
 func LoadFile(filename string, enc Encoding) (*Properties, error) {
 	l := &Loader{Encoding: enc}
diff --git a/parser.go b/parser.go
index cdc4a80..6cc3865 100644
--- a/parser.go
+++ b/parser.go
@@ -7,30 +7,53 @@ package properties
 import (
 	"fmt"
 	"runtime"
+	"strings"
 )
 
 type parser struct {
 	lex *lexer
 }
 
-func parse(input string) (properties *Properties, err error) {
-	p := &parser{lex: lex(input)}
+func parse(input string, preserveFormatting bool) (properties *Properties, err error) {
+	p := &parser{lex: lex(input, preserveFormatting)}
 	defer p.recover(&err)
 
 	properties = NewProperties()
+	properties.PreserveFormatting = preserveFormatting
 	key := ""
-	comments := []string{}
+	comments := []Comment{}
 
 	for {
 		token := p.expectOneOf(itemComment, itemKey, itemEOF)
 		switch token.typ {
 		case itemEOF:
+			if preserveFormatting && (len(comments) > 0 || token.val != "") {
+				// There are comments at the end of the input that are not tied to a particular key
+				// Save these off against a special empty key when preserving formatting
+				if token.val != "" {
+					prefixIndex := 0
+					// Include leading whitespace into the prefix
+					prefixIndex = strings.Index(token.val, strings.TrimSpace(token.val))
+					prefix := token.val[0 : prefixIndex+1]
+					comment := Comment{prefix, token.val[prefixIndex+1 : len(token.val)]}
+					comments = append(comments, comment)
+				}
+				properties.c[""] = comments
+			}
 			goto done
 		case itemComment:
-			comments = append(comments, token.val)
+			prefix := "#"
+			prefixIndex := 0
+			if preserveFormatting {
+				// Include leading whitespace into the prefix
+				prefixIndex = strings.Index(token.val, strings.TrimSpace(token.val))
+				prefix = token.val[0:prefixIndex+1]
+			}
+			comment := Comment{prefix, token.val[prefixIndex+1:len(token.val)]}
+			comments = append(comments, comment)
 			continue
 		case itemKey:
-			key = token.val
+			key = strings.TrimSpace(token.val)
 			if _, ok := properties.m[key]; !ok {
 				properties.k = append(properties.k, key)
 			}
@@ -39,7 +62,7 @@ func parse(input string) (properties *Properties, err error) {
 		token = p.expectOneOf(itemValue, itemEOF)
 		if len(comments) > 0 {
 			properties.c[key] = comments
-			comments = []string{}
+			comments = []Comment{}
 		}
 		switch token.typ {
 		case itemEOF:
diff --git a/properties.go b/properties.go
index cb3d1a3..0776b5e 100644
--- a/properties.go
+++ b/properties.go
@@ -48,6 +48,11 @@ func PanicHandler(err error) {
 
 // -----------------------------------------------------------------------------
 
+type Comment struct {
+	prefix string // prefix type (#, !, or "" to suppress a prefix)
+	val string   // The value of this item.
+}
+
 // A Properties contains the key/value pairs from the properties input.
 // All values are stored in unexpanded form and are expanded at runtime
 type Properties struct {
@@ -65,10 +70,13 @@ type Properties struct {
 	m map[string]string
 
 	// Stores the comments per key.
-	c map[string][]string
+	c map[string][]Comment
 
 	// Stores the keys in order of appearance.
 	k []string
+
+	// Preserve formatting (whitespace/comment prefix) when reading in properties
+	PreserveFormatting bool
 }
 
 // NewProperties creates a new Properties struct with the default
@@ -78,7 +86,7 @@ func NewProperties() *Properties {
 		Prefix:  "${",
 		Postfix: "}",
 		m:       map[string]string{},
-		c:       map[string][]string{},
+		c:       map[string][]Comment{},
 		k:       []string{},
 	}
 }
@@ -131,7 +139,7 @@ func (p *Properties) MustGet(key string) string {
 
 // ClearComments removes the comments for all keys.
 func (p *Properties) ClearComments() {
-	p.c = map[string][]string{}
+	p.c = map[string][]Comment{}
 }
 
 // ----------------------------------------------------------------------------
@@ -142,7 +150,7 @@ func (p *Properties) GetComment(key string) string {
 	if !ok || len(comments) == 0 {
 		return ""
 	}
-	return comments[len(comments)-1]
+	return comments[len(comments)-1].val
 }
 
 // ----------------------------------------------------------------------------
@@ -150,7 +158,11 @@ func (p *Properties) GetComment(key string) string {
 // GetComments returns all comments that appeared before the given key or nil.
 func (p *Properties) GetComments(key string) []string {
 	if comments, ok := p.c[key]; ok {
-		return comments
+		var list []string
+		for _, comment := range comments {
+			list = append(list, comment.val)
+		}
+		return list
 	}
 	return nil
 }
@@ -159,7 +171,14 @@ func (p *Properties) GetComments(key string) []string {
 
 // SetComment sets the comment for the key.
 func (p *Properties) SetComment(key, comment string) {
-	p.c[key] = []string{comment}
+	p.c[key] = []Comment{{"#", comment}}
+}
+
+// ----------------------------------------------------------------------------
+
+// SetComment sets the comment for the key.
+func (p *Properties) SetCommentWithPrefix(key, prefix string, comment string) {
+	p.c[key] = []Comment{{prefix, comment}}
 }
 
 // ----------------------------------------------------------------------------
@@ -167,11 +186,23 @@ func (p *Properties) SetComment(key, comment string) {
 // SetComments sets the comments for the key. If the comments are nil then
 // all comments for this key are deleted.
 func (p *Properties) SetComments(key string, comments []string) {
+	p.SetCommentsWithPrefix(key, "#", comments)
+}
+
+// ----------------------------------------------------------------------------
+
+// SetComments sets the comments for the key. If the comments are nil then
+// all comments for this key are deleted.
+func (p *Properties) SetCommentsWithPrefix(key string, prefix string, comments []string) {
 	if comments == nil {
 		delete(p.c, key)
 		return
 	}
-	p.c[key] = comments
+	var list = make([]Comment, 0, len(comments))
+	for _, comment := range comments {
+		list = append(list, Comment{prefix, comment})
+	}
+	p.c[key] = list
 }
 
 // ----------------------------------------------------------------------------
@@ -592,30 +623,33 @@ func (p *Properties) Write(w io.Writer, enc Encoding) (n int, err error) {
 	return p.WriteComment(w, "", enc)
 }
 
-// WriteComment writes all unexpanced 'key = value' pairs to the given writer.
+// WriteComment writes all unexpanded 'key = value' pairs to the given writer.
 // If prefix is not empty then comments are written with a blank line and the
-// given prefix. The prefix should be either "# " or "! " to be compatible with
+// given prefix. The prefix should be either "#" or "!" to be compatible with
 // the properties file format. Otherwise, the properties parser will not be
 // able to read the file back in. It returns the number of bytes written and
 // any write error encountered.
 func (p *Properties) WriteComment(w io.Writer, prefix string, enc Encoding) (n int, err error) {
-	var x int
+	return p.WriteCommentWithFormatting(w, prefix, enc, false)
+}
 
+func (p *Properties) WriteCommentWithFormatting(w io.Writer, prefix string, enc Encoding, preserveFormatting bool) (n int, err error) {
+	var x int
 	for _, key := range p.k {
 		value := p.m[key]
 
-		if prefix != "" {
+		if prefix != "" || preserveFormatting {
 			if comments, ok := p.c[key]; ok {
 				// don't print comments if they are all empty
 				allEmpty := true
 				for _, c := range comments {
-					if c != "" {
+					if strings.TrimSpace(c.val) != "" {
 						allEmpty = false
 						break
 					}
 				}
 
-				if !allEmpty {
+				if !allEmpty && !preserveFormatting {
 					// add a blank line between entries but not at the top
 					if len(comments) > 0 && n > 0 {
 						x, err = fmt.Fprintln(w)
@@ -624,13 +658,25 @@ func (p *Properties) WriteComment(w io.Writer, prefix string, enc Encoding) (n i
 						}
 						n += x
 					}
-
+				}
+				if !allEmpty || preserveFormatting {
 					for _, c := range comments {
-						x, err = fmt.Fprintf(w, "%s%s\n", prefix, encode(c, "", enc))
-						if err != nil {
-							return
+						actualPrefix := prefix
+						if prefix == "" && c.prefix != "\n" {
+							// No prefix passed in and this is not just a newline comment, so use the stored prefix
+							actualPrefix = c.prefix
+						}
+						actualValue := c.val
+						if !preserveFormatting {
+							actualValue = strings.TrimSpace(c.val)
+						}
+						if preserveFormatting || strings.TrimSpace(actualValue) != "" {
+							x, err = fmt.Fprintf(w, "%s%s\n", actualPrefix, encode(actualValue, "", enc))
+							if err != nil {
+								return
+							}
+							n += x
 						}
-						n += x
 					}
 				}
 			}
@@ -642,6 +688,21 @@ func (p *Properties) WriteComment(w io.Writer, prefix string, enc Encoding) (n i
 		}
 		n += x
 	}
+	if comments, exists := p.c[""]; exists && preserveFormatting {
+		for _, c := range comments {
+			actualPrefix := prefix
+			if prefix == "" && c.prefix != "\n" {
+				// No prefix passed in and this is not just a newline comment, so use the stored prefix
+				actualPrefix = c.prefix
+			}
+			actualValue := c.val
+			x, err = fmt.Fprintf(w, "%s%s\n", actualPrefix, encode(actualValue, "", enc))
+			if err != nil {
+				return
+			}
+			n += x
+		}
+	}
 	return
 }
 
diff --git a/properties_test.go b/properties_test.go
index 5db64c3..882a4f6 100644
--- a/properties_test.go
+++ b/properties_test.go
@@ -787,7 +787,7 @@ func TestMustSet(t *testing.T) {
 
 func TestWrite(t *testing.T) {
 	for _, test := range writeTests {
-		p, err := parse(test.input)
+		p, err := parse(test.input, false)
 
 		buf := new(bytes.Buffer)
 		var n int
@@ -806,7 +806,7 @@ func TestWrite(t *testing.T) {
 
 func TestWriteComment(t *testing.T) {
 	for _, test := range writeCommentTests {
-		p, err := parse(test.input)
+		p, err := parse(test.input, false)
 
 		buf := new(bytes.Buffer)
 		var n int
@@ -970,7 +970,7 @@ func assertKeyValues(t *testing.T, input string, p *Properties, keyvalues ...str
 }
 
 func mustParse(t *testing.T, s string) *Properties {
-	p, err := parse(s)
+	p, err := parse(s, false)
 	if err != nil {
 		t.Fatalf("parse failed with %s", err)
 	}

From 9db77aeb1167f61f08ef6aecba243478ad7b4e4d Mon Sep 17 00:00:00 2001
From: Arvind Thirunarayanan <athirunar@confluent.io>
Date: Tue, 30 Jul 2019 10:45:00 -0500
Subject: [PATCH 2/5] Add tests for preserveFormatting feature

---
 properties_test.go | 135 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 135 insertions(+)

diff --git a/properties_test.go b/properties_test.go
index 882a4f6..b483164 100644
--- a/properties_test.go
+++ b/properties_test.go
@@ -122,6 +122,24 @@ var commentTests = []struct {
 
 // ----------------------------------------------------------------------------
 
+var formattedCommentTests = []struct {
+	prefix string
+	comments []string
+}{
+	{"", nil},
+	{"", []string{""}},
+	{" ", []string{"   "}},
+	{"#", []string{"comment"}},
+	{"!", []string{"comment"}},
+	{"#", []string{""}},
+	{"!", []string{""}},
+	{"\n", []string{""}},
+	{"  #", []string{"comment1", "comment2"}},
+	{" !", []string{"comment1", "comment2"}},
+}
+
+// ----------------------------------------------------------------------------
+
 var errorTests = []struct {
 	input, msg string
 }{
@@ -173,6 +191,62 @@ var writeCommentTests = []struct {
 	{"# comment\n\nkey = value", "# comment\nkey = value\n", "ISO-8859-1"},
 	{"# comment1\n# comment2\nkey = value", "# comment1\n# comment2\nkey = value\n", "ISO-8859-1"},
 	{"#comment1\nkey1 = value1\n#comment2\nkey2 = value2", "# comment1\nkey1 = value1\n\n# comment2\nkey2 = value2\n", "ISO-8859-1"},
+	{"key = value\n # comment1", "key = value\n", "ISO-8859-1"},
+
+	// UTF-8 tests
+	{"key = value", "key = value\n", "UTF-8"},
+	{"# comment⌘\nkey = value⌘", "# comment⌘\nkey = value⌘\n", "UTF-8"},
+	{"\n# comment⌘\nkey = value⌘", "# comment⌘\nkey = value⌘\n", "UTF-8"},
+	{"# comment⌘\n\nkey = value⌘", "# comment⌘\nkey = value⌘\n", "UTF-8"},
+	{"# comment1⌘\n# comment2⌘\nkey = value⌘", "# comment1⌘\n# comment2⌘\nkey = value⌘\n", "UTF-8"},
+	{"#comment1⌘\nkey1 = value1⌘\n#comment2⌘\nkey2 = value2⌘", "# comment1⌘\nkey1 = value1⌘\n\n# comment2⌘\nkey2 = value2⌘\n", "UTF-8"},
+}
+
+// ----------------------------------------------------------------------------
+
+var writeCommentWithFormattingTests = []struct {
+	input, encoding string
+}{
+	// ISO-8859-1 tests
+	{"key = value", "ISO-8859-1"},
+	{"#\nkey = value","ISO-8859-1"},
+	{"#\n#\n#\nkey = value", "ISO-8859-1"},
+	{"# comment\nkey = value", "ISO-8859-1"},
+	{"\n# comment\nkey = value", "ISO-8859-1"},
+	{"# comment\n\nkey = value", "ISO-8859-1"},
+	{"# comment1\n# comment2\nkey = value", "ISO-8859-1"},
+	{"#comment1\nkey1 = value1\n #comment2\nkey2 = value2", "ISO-8859-1"},
+	{"\n\n # \n#\n#comment\n\nkey = value", "ISO-8859-1"},
+	{"\n#comment1\nkey1 = value1\n\n#\n#comment2#\n\nkey2 = value2", "ISO-8859-1"},
+	{"key = value\n # comment1", "ISO-8859-1"},
+
+	// UTF-8 tests
+	{"key = value", "UTF-8"},
+	{"# comment⌘\nkey = value⌘", "UTF-8"},
+	{"\n# comment⌘\nkey = value⌘", "UTF-8"},
+	{"# comment⌘\n\nkey = value⌘", "UTF-8"},
+	{"# comment1⌘\n# comment2⌘\nkey = value⌘", "UTF-8"},
+	{"#comment1⌘\nkey1 = value1⌘\n#comment2⌘\nkey2 = value2⌘", "UTF-8"},
+	{"\n\n # \n#\n#comment\n\nkey = value⌘", "UTF-8"},
+}
+
+// ----------------------------------------------------------------------------
+
+var writeFormattedCommentWithoutFormattingTests = []struct {
+	input, output, encoding string
+}{
+	// ISO-8859-1 tests
+	{"key = value", "key = value\n", "ISO-8859-1"},
+	{"#\nkey = value", "key = value\n", "ISO-8859-1"},
+	{"#\n#\n#\nkey = value", "key = value\n", "ISO-8859-1"},
+	{" # comment\nkey = value", "# comment\nkey = value\n", "ISO-8859-1"},
+	{"\n# comment\nkey = value", "# comment\nkey = value\n", "ISO-8859-1"},
+	{"# comment\n\nkey = value", "# comment\nkey = value\n", "ISO-8859-1"},
+	{"# comment1\n# comment2\nkey = value", "# comment1\n# comment2\nkey = value\n", "ISO-8859-1"},
+	{"#comment1\nkey1 = value1\n #comment2\nkey2 = value2", "# comment1\nkey1 = value1\n\n# comment2\nkey2 = value2\n", "ISO-8859-1"},
+	{"\n\n # \n#\n #comment\n\nkey = value", "# comment\nkey = value\n","ISO-8859-1"},
+	{"\n#comment1\nkey1 = value1\n\n#\n#comment2\n#\n\nkey2 = value2", "# comment1\nkey1 = value1\n\n# comment2\nkey2 = value2\n", "ISO-8859-1"},
+	{"key = value\n # comment1", "key = value\n", "ISO-8859-1"},
 
 	// UTF-8 tests
 	{"key = value", "key = value\n", "UTF-8"},
@@ -683,6 +757,28 @@ func TestComment(t *testing.T) {
 	}
 }
 
+func TestFormattedComment(t *testing.T) {
+	for _, test := range formattedCommentTests {
+		key := "key"
+		p := LoadMap(map[string]string{key: "value"})
+
+		// test setting comments
+		if len(test.comments) > 0 {
+			// set single comment
+			p.ClearComments()
+			assert.Equal(t, len(p.c), 0)
+			p.SetCommentWithPrefix(key, test.prefix, test.comments[0])
+			assert.Equal(t, p.GetComment(key), test.comments[0])
+
+			// set multiple comments
+			p.ClearComments()
+			assert.Equal(t, len(p.c), 0)
+			p.SetCommentsWithPrefix(key, test.prefix, test.comments)
+			assert.Equal(t, p.GetComments(key), test.comments)
+		}
+	}
+}
+
 func TestFilter(t *testing.T) {
 	for _, test := range filterTests {
 		p := mustParse(t, test.input)
@@ -823,6 +919,45 @@ func TestWriteComment(t *testing.T) {
 	}
 }
 
+func TestWriteCommentWithFormatting(t *testing.T) {
+	for _, test := range writeCommentWithFormattingTests {
+		p, err := parse(test.input, true)
+
+		buf := new(bytes.Buffer)
+		var n int
+		switch test.encoding {
+		case "UTF-8":
+			n, err = p.WriteCommentWithFormatting(buf, "", UTF8, true)
+		case "ISO-8859-1":
+			n, err = p.WriteCommentWithFormatting(buf, "", ISO_8859_1, true)
+		}
+		assert.Equal(t, err, nil)
+		s := string(buf.Bytes())
+		expectedOutput := test.input + "\n"	// output should match input with trailing newline
+		assert.Equal(t, n, len(expectedOutput), fmt.Sprintf("input=%q expected=%q obtained=%q", test.input, expectedOutput, s))
+		assert.Equal(t, s, expectedOutput, fmt.Sprintf("input=%q expected=%q obtained=%q", test.input, expectedOutput, s))
+	}
+}
+
+func TestWriteFormattedCommentWithoutFormatting(t *testing.T) {
+	for _, test := range writeFormattedCommentWithoutFormattingTests {
+		p, err := parse(test.input, true)
+
+		buf := new(bytes.Buffer)
+		var n int
+		switch test.encoding {
+		case "UTF-8":
+			n, err = p.WriteComment(buf, "# ", UTF8)
+		case "ISO-8859-1":
+			n, err = p.WriteComment(buf, "# ", ISO_8859_1)
+		}
+		assert.Equal(t, err, nil)
+		s := string(buf.Bytes())
+		assert.Equal(t, n, len(test.output), fmt.Sprintf("input=%q expected=%q obtained=%q", test.input, test.output, s))
+		assert.Equal(t, s, test.output, fmt.Sprintf("input=%q expected=%q obtained=%q", test.input, test.output, s))
+	}
+}
+
 func TestCustomExpansionExpression(t *testing.T) {
 	testKeyValuePrePostfix(t, "*[", "]*", "key=value\nkey2=*[key]*", "key", "value", "key2", "value")
 }

From 4ddf0d00aa931edb74099b197002f2de8ec177ad Mon Sep 17 00:00:00 2001
From: Arvind Thirunarayanan <athirunar@confluent.io>
Date: Tue, 6 Aug 2019 14:31:51 -0500
Subject: [PATCH 3/5] PR comments

---
 lex.go        | 62 ++++++++++++++++++++++++---------------------------
 load.go       | 14 +++++-------
 parser.go     | 41 ++++++++++++++++++++--------------
 properties.go | 24 ++++++++++----------
 4 files changed, 70 insertions(+), 71 deletions(-)

diff --git a/lex.go b/lex.go
index 7f376c3..33251b3 100644
--- a/lex.go
+++ b/lex.go
@@ -60,7 +60,7 @@ type stateFn func(*lexer) stateFn
 
 // lexer holds the state of the scanner.
 type lexer struct {
-	preserveFormatting bool  // whether to scan EOLs/whitespace as part of comments
+	keepWS bool       // retain EOLs/whitespace as part of comments
 	input   string    // the string being scanned
 	state   stateFn   // the next lexing function to enter
 	pos     int       // current position in the input
@@ -163,15 +163,12 @@ func (l *lexer) nextItem() item {
 }
 
 // lex creates a new scanner for the input string.
-func lex(input string, preserveFormatting bool) *lexer {
-	l := &lexer{
-		preserveFormatting: preserveFormatting,
+func lex(input string) *lexer {
+	return &lexer{
 		input: input,
 		items: make(chan item),
 		runes: make([]rune, 0, 32),
 	}
-	go l.run()
-	return l
 }
 
 // run runs the state machine for the lexer.
@@ -190,37 +187,36 @@ func lexBeforeKey(l *lexer) stateFn {
 		l.emit(itemEOF)
 		return nil
 
+	case isEOL(r) && l.keepWS:
+		// treat as part of the next key's comments block
+		// add the EOL rune r as the comment's prefix
+		l.appendRune(r)
+		l.backup()
+		return lexComment
+
 	case isEOL(r):
-		if l.preserveFormatting {
-			// treat as part of the next key's comments block
-			// add the EOL rune r as the comment's prefix
-			l.appendRune(r)
-			l.backup()
-			return lexComment
-		} else {
-			l.ignore()
-			return lexBeforeKey
-		}
+		l.ignore()
+		return lexBeforeKey
+
+	case isComment(r) && l.keepWS:
+		// treat as part of the next key's comments block
+		// use existing prefix r as the comment's prefix
+		l.appendRune(r)
+		return lexComment
 
 	case isComment(r):
-		if l.preserveFormatting {
-			// treat as part of the next key's comments block
-			// use existing prefix r as the comment's prefix
-			l.appendRune(r)
-		} else {
-			// use "#" as the comment's prefix
-			l.appendRune('#')
-		}
+		// use "#" as the comment's prefix
+		l.appendRune('#')
 		return lexComment
 
+	case isWhitespace(r) && l.keepWS:
+		// treat as part of the next key's comments block
+		// add the whitespace rune r as the comment's prefix
+		l.appendRune(r)
+		return lexBeforeKey
+
 	case isWhitespace(r):
-		if l.preserveFormatting {
-			// treat as part of the next key's comments block
-			// add the whitespace rune r as the comment's prefix
-			l.appendRune(r)
-		} else {
-			l.ignore()
-		}
+		l.ignore()
 		return lexBeforeKey
 
 	default:
@@ -231,7 +227,7 @@ func lexBeforeKey(l *lexer) stateFn {
 
 // lexComment scans a comment line. The comment character has already been scanned.
 func lexComment(l *lexer) stateFn {
-	if ! l.preserveFormatting {
+	if ! l.keepWS {
 		l.acceptRun(whitespace)
 		l.ignore()
 	}
@@ -239,7 +235,7 @@ func lexComment(l *lexer) stateFn {
 
 		switch r := l.next(); {
 		case isEOF(r):
-			if ! l.preserveFormatting {
+			if ! l.keepWS {
 				l.ignore()
 			}
 			l.emit(itemEOF)
diff --git a/load.go b/load.go
index 11250de..90bdbf0 100644
--- a/load.go
+++ b/load.go
@@ -54,7 +54,7 @@ type Loader struct {
 
 // Load reads a buffer into a Properties struct.
 func (l *Loader) LoadBytes(buf []byte) (*Properties, error) {
-	return l.loadBytes(buf, l.Encoding, l.PreserveFormatting)
+	return l.loadBytes(buf, l.Encoding)
 }
 
 // LoadAll reads the content of multiple URLs or files in the given order into
@@ -104,7 +104,7 @@ func (l *Loader) LoadFile(filename string) (*Properties, error) {
 		}
 		return nil, err
 	}
-	return l.loadBytes(data, l.Encoding, l.PreserveFormatting)
+	return l.loadBytes(data, l.Encoding)
 }
 
 // LoadURL reads the content of the URL into a Properties struct.
@@ -147,11 +147,11 @@ func (l *Loader) LoadURL(url string) (*Properties, error) {
 		return nil, fmt.Errorf("properties: invalid content type %s", ct)
 	}
 
-	return l.loadBytes(body, enc, l.PreserveFormatting)
+	return l.loadBytes(body, enc)
 }
 
-func (l *Loader) loadBytes(buf []byte, enc Encoding, preserveFormatting bool) (*Properties, error) {
-	p, err := parse(convert(buf, enc), preserveFormatting)
+func (l *Loader) loadBytes(buf []byte, enc Encoding) (*Properties, error) {
+	p, err := parse(convert(buf, enc), l.PreserveFormatting)
 	if err != nil {
 		return nil, err
 	}
@@ -183,10 +183,6 @@ func LoadMap(m map[string]string) *Properties {
 	return p
 }
 
-func GetLoader() (*Loader, error) {
-	return &Loader{}, nil
-}
-
 // LoadFile reads a file into a Properties struct.
 func LoadFile(filename string, enc Encoding) (*Properties, error) {
 	l := &Loader{Encoding: enc}
diff --git a/parser.go b/parser.go
index 6cc3865..757b3fd 100644
--- a/parser.go
+++ b/parser.go
@@ -15,31 +15,38 @@ type parser struct {
 }
 
 func parse(input string, preserveFormatting bool) (properties *Properties, err error) {
-	p := &parser{lex: lex(input, preserveFormatting)}
+	l := lex(input)
+	l.keepWS = preserveFormatting
+	go l.run()
+	p := &parser{lex: l}
 	defer p.recover(&err)
 
 	properties = NewProperties()
-	properties.PreserveFormatting = preserveFormatting
 	key := ""
-	comments := []Comment{}
+	comments := []prefixedComment{}
 
 	for {
 		token := p.expectOneOf(itemComment, itemKey, itemEOF)
 		switch token.typ {
 		case itemEOF:
-			if preserveFormatting && (len(comments) > 0 || token.val != "") {
-				// There are comments at the end of the input that are not tied to a particular key
-				// Save these off against a special empty key when preserving formatting
-				if token.val != "" {
-					prefixIndex := 0
-					// Include leading whitespace into the prefix
-					prefixIndex = strings.Index(token.val, strings.TrimSpace(token.val))
-					prefix := token.val[0 : prefixIndex+1]
-					comment := Comment{prefix, token.val[prefixIndex+1 : len(token.val)]}
-					comments = append(comments, comment)
-				}
-				properties.c[""] = comments
+			if !preserveFormatting || (len(comments) == 0 && token.val == "") {
+				goto done
 			}
+			// There are comments at the end of the input that are not tied to a particular key
+			// Save these off when preserving formatting
+			if token.val == "" {
+				// Save off previously parsed trailing comments
+				properties.trailingComments = comments
+				goto done
+			}
+			// Need to save off the last comment line as well
+			prefixIndex := 0
+			// Include leading whitespace into the prefix
+			prefixIndex = strings.Index(token.val, strings.TrimSpace(token.val))
+			prefix := token.val[0 : prefixIndex+1]
+			comment := prefixedComment{prefix, token.val[prefixIndex+1 : len(token.val)]}
+			comments = append(comments, comment)
+			properties.trailingComments = comments
 			goto done
 		case itemComment:
 			prefix := "#"
@@ -49,7 +56,7 @@ func parse(input string, preserveFormatting bool) (properties *Properties, err e
 				prefixIndex = strings.Index(token.val, strings.TrimSpace(token.val))
 				prefix = token.val[0:prefixIndex+1]
 			}
-			comment := Comment{prefix, token.val[prefixIndex+1:len(token.val)]}
+			comment := prefixedComment{prefix, token.val[prefixIndex+1:len(token.val)]}
 			comments = append(comments, comment)
 			continue
 		case itemKey:
@@ -62,7 +69,7 @@ func parse(input string, preserveFormatting bool) (properties *Properties, err e
 		token = p.expectOneOf(itemValue, itemEOF)
 		if len(comments) > 0 {
 			properties.c[key] = comments
-			comments = []Comment{}
+			comments = []prefixedComment{}
 		}
 		switch token.typ {
 		case itemEOF:
diff --git a/properties.go b/properties.go
index 0776b5e..09d7ee1 100644
--- a/properties.go
+++ b/properties.go
@@ -48,7 +48,7 @@ func PanicHandler(err error) {
 
 // -----------------------------------------------------------------------------
 
-type Comment struct {
+type prefixedComment struct {
 	prefix string // prefix type (#, !, or "" to suppress a prefix)
 	val string   // The value of this item.
 }
@@ -70,13 +70,13 @@ type Properties struct {
 	m map[string]string
 
 	// Stores the comments per key.
-	c map[string][]Comment
+	c map[string][]prefixedComment
 
 	// Stores the keys in order of appearance.
 	k []string
 
-	// Preserve formatting (whitespace/comment prefix) when reading in properties
-	PreserveFormatting bool
+	// Stores any trailing comments if we want to preserveFormatting
+	trailingComments []prefixedComment
 }
 
 // NewProperties creates a new Properties struct with the default
@@ -86,7 +86,7 @@ func NewProperties() *Properties {
 		Prefix:  "${",
 		Postfix: "}",
 		m:       map[string]string{},
-		c:       map[string][]Comment{},
+		c:       map[string][]prefixedComment{},
 		k:       []string{},
 	}
 }
@@ -139,7 +139,7 @@ func (p *Properties) MustGet(key string) string {
 
 // ClearComments removes the comments for all keys.
 func (p *Properties) ClearComments() {
-	p.c = map[string][]Comment{}
+	p.c = map[string][]prefixedComment{}
 }
 
 // ----------------------------------------------------------------------------
@@ -171,14 +171,14 @@ func (p *Properties) GetComments(key string) []string {
 
 // SetComment sets the comment for the key.
 func (p *Properties) SetComment(key, comment string) {
-	p.c[key] = []Comment{{"#", comment}}
+	p.c[key] = []prefixedComment{{"#", comment}}
 }
 
 // ----------------------------------------------------------------------------
 
 // SetComment sets the comment for the key.
 func (p *Properties) SetCommentWithPrefix(key, prefix string, comment string) {
-	p.c[key] = []Comment{{prefix, comment}}
+	p.c[key] = []prefixedComment{{prefix, comment}}
 }
 
 // ----------------------------------------------------------------------------
@@ -198,9 +198,9 @@ func (p *Properties) SetCommentsWithPrefix(key string, prefix string, comments [
 		delete(p.c, key)
 		return
 	}
-	var list = make([]Comment, 0, len(comments))
+	var list = make([]prefixedComment, 0, len(comments))
 	for _, comment := range comments {
-		list = append(list, Comment{prefix, comment})
+		list = append(list, prefixedComment{prefix, comment})
 	}
 	p.c[key] = list
 }
@@ -688,8 +688,8 @@ func (p *Properties) WriteCommentWithFormatting(w io.Writer, prefix string, enc
 		}
 		n += x
 	}
-	if comments, exists := p.c[""]; exists && preserveFormatting {
-		for _, c := range comments {
+	if preserveFormatting {
+		for _, c := range p.trailingComments {
 			actualPrefix := prefix
 			if prefix == "" && c.prefix != "\n" {
 				// No prefix passed in and this is not just a newline comment, so use the stored prefix

From b0894f28c51c0ae18f40ad3a07df61848656b6f1 Mon Sep 17 00:00:00 2001
From: Arvind Thirunarayanan <athirunar@confluent.io>
Date: Tue, 6 Aug 2019 18:07:13 -0500
Subject: [PATCH 4/5] Create separate WriteFormattedContent method to simplify
 properties.go

  - Original behavior is still retained for calls to the original Write methods.
---
 lex.go        | 10 ++-----
 properties.go | 83 ++++++++++++++++++++++++++++++---------------------
 2 files changed, 51 insertions(+), 42 deletions(-)

diff --git a/lex.go b/lex.go
index 33251b3..5da510c 100644
--- a/lex.go
+++ b/lex.go
@@ -198,20 +198,15 @@ func lexBeforeKey(l *lexer) stateFn {
 		l.ignore()
 		return lexBeforeKey
 
-	case isComment(r) && l.keepWS:
+	case isComment(r):
 		// treat as part of the next key's comments block
 		// use existing prefix r as the comment's prefix
 		l.appendRune(r)
 		return lexComment
 
-	case isComment(r):
-		// use "#" as the comment's prefix
-		l.appendRune('#')
-		return lexComment
-
 	case isWhitespace(r) && l.keepWS:
 		// treat as part of the next key's comments block
-		// add the whitespace rune r as the comment's prefix
+		// add the whitespace rune r as part of the comment's prefix
 		l.appendRune(r)
 		return lexBeforeKey
 
@@ -232,7 +227,6 @@ func lexComment(l *lexer) stateFn {
 		l.ignore()
 	}
 	for {
-
 		switch r := l.next(); {
 		case isEOF(r):
 			if ! l.keepWS {
diff --git a/properties.go b/properties.go
index 09d7ee1..0a59091 100644
--- a/properties.go
+++ b/properties.go
@@ -49,8 +49,8 @@ func PanicHandler(err error) {
 // -----------------------------------------------------------------------------
 
 type prefixedComment struct {
-	prefix string // prefix type (#, !, or "" to suppress a prefix)
-	val string   // The value of this item.
+	prefix string // prefix type (#, !)
+	val string    // The comment string after the prefix
 }
 
 // A Properties contains the key/value pairs from the properties input.
@@ -625,20 +625,17 @@ func (p *Properties) Write(w io.Writer, enc Encoding) (n int, err error) {
 
 // WriteComment writes all unexpanded 'key = value' pairs to the given writer.
 // If prefix is not empty then comments are written with a blank line and the
-// given prefix. The prefix should be either "#" or "!" to be compatible with
+// given prefix. The prefix should be either "# " or "! " to be compatible with
 // the properties file format. Otherwise, the properties parser will not be
 // able to read the file back in. It returns the number of bytes written and
 // any write error encountered.
 func (p *Properties) WriteComment(w io.Writer, prefix string, enc Encoding) (n int, err error) {
-	return p.WriteCommentWithFormatting(w, prefix, enc, false)
-}
-
-func (p *Properties) WriteCommentWithFormatting(w io.Writer, prefix string, enc Encoding, preserveFormatting bool) (n int, err error) {
 	var x int
+
 	for _, key := range p.k {
 		value := p.m[key]
 
-		if prefix != "" || preserveFormatting {
+		if prefix != "" {
 			if comments, ok := p.c[key]; ok {
 				// don't print comments if they are all empty
 				allEmpty := true
@@ -649,7 +646,7 @@ func (p *Properties) WriteCommentWithFormatting(w io.Writer, prefix string, enc
 					}
 				}
 
-				if !allEmpty && !preserveFormatting {
+				if !allEmpty {
 					// add a blank line between entries but not at the top
 					if len(comments) > 0 && n > 0 {
 						x, err = fmt.Fprintln(w)
@@ -658,20 +655,11 @@ func (p *Properties) WriteCommentWithFormatting(w io.Writer, prefix string, enc
 						}
 						n += x
 					}
-				}
-				if !allEmpty || preserveFormatting {
+
 					for _, c := range comments {
-						actualPrefix := prefix
-						if prefix == "" && c.prefix != "\n" {
-							// No prefix passed in and this is not just a newline comment, so use the stored prefix
-							actualPrefix = c.prefix
-						}
-						actualValue := c.val
-						if !preserveFormatting {
-							actualValue = strings.TrimSpace(c.val)
-						}
-						if preserveFormatting || strings.TrimSpace(actualValue) != "" {
-							x, err = fmt.Fprintf(w, "%s%s\n", actualPrefix, encode(actualValue, "", enc))
+						comment := strings.TrimSpace(c.val)
+						if comment != "" {
+							x, err = fmt.Fprintf(w, "%s%s\n", prefix, encode(comment, "", enc))
 							if err != nil {
 								return
 							}
@@ -688,20 +676,47 @@ func (p *Properties) WriteCommentWithFormatting(w io.Writer, prefix string, enc
 		}
 		n += x
 	}
-	if preserveFormatting {
-		for _, c := range p.trailingComments {
-			actualPrefix := prefix
-			if prefix == "" && c.prefix != "\n" {
-				// No prefix passed in and this is not just a newline comment, so use the stored prefix
-				actualPrefix = c.prefix
-			}
-			actualValue := c.val
-			x, err = fmt.Fprintf(w, "%s%s\n", actualPrefix, encode(actualValue, "", enc))
-			if err != nil {
-				return
+	return
+}
+
+// WriteFormattedComment writes all unexpanded 'key = value' pairs to the given writer.
+// Comments are written out with their original preserved formatting. It returns the number of bytes
+// written and any write error encountered.
+func (p *Properties) WriteFormattedComment(w io.Writer, enc Encoding) (n int, err error) {
+	var x int
+	for _, key := range p.k {
+		value := p.m[key]
+
+		if comments, ok := p.c[key]; ok {
+			for _, c := range comments {
+				prefix := ""
+				if (c.prefix != "\n") {
+					prefix = c.prefix
+				}
+				x, err = fmt.Fprintf(w, "%s%s\n", prefix, encode(c.val, "", enc))
+				if err != nil {
+					return
+				}
+				n += x
 			}
-			n += x
 		}
+
+		x, err = fmt.Fprintf(w, "%s = %s\n", encode(key, " :", enc), encode(value, "", enc))
+		if err != nil {
+			return
+		}
+		n += x
+	}
+	for _, c := range p.trailingComments {
+		prefix := ""
+		if (c.prefix != "\n") {
+			prefix = c.prefix
+		}
+		x, err = fmt.Fprintf(w, "%s%s\n", prefix, encode(c.val, "", enc))
+		if err != nil {
+			return
+		}
+		n += x
 	}
 	return
 }

From c94c1ece9b37a28dcf37bfd94ec24dd24c09f723 Mon Sep 17 00:00:00 2001
From: Arvind Thirunarayanan <athirunar@confluent.io>
Date: Tue, 6 Aug 2019 18:14:13 -0500
Subject: [PATCH 5/5] Test updates for new WriteFormattedComment method

---
 properties_test.go | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/properties_test.go b/properties_test.go
index b483164..d2b11ee 100644
--- a/properties_test.go
+++ b/properties_test.go
@@ -884,6 +884,7 @@ func TestMustSet(t *testing.T) {
 func TestWrite(t *testing.T) {
 	for _, test := range writeTests {
 		p, err := parse(test.input, false)
+		assert.Equal(t, err, nil)
 
 		buf := new(bytes.Buffer)
 		var n int
@@ -903,6 +904,7 @@ func TestWrite(t *testing.T) {
 func TestWriteComment(t *testing.T) {
 	for _, test := range writeCommentTests {
 		p, err := parse(test.input, false)
+		assert.Equal(t, err, nil)
 
 		buf := new(bytes.Buffer)
 		var n int
@@ -922,14 +924,15 @@ func TestWriteComment(t *testing.T) {
 func TestWriteCommentWithFormatting(t *testing.T) {
 	for _, test := range writeCommentWithFormattingTests {
 		p, err := parse(test.input, true)
+		assert.Equal(t, err, nil)
 
 		buf := new(bytes.Buffer)
 		var n int
 		switch test.encoding {
 		case "UTF-8":
-			n, err = p.WriteCommentWithFormatting(buf, "", UTF8, true)
+			n, err = p.WriteFormattedComment(buf, UTF8)
 		case "ISO-8859-1":
-			n, err = p.WriteCommentWithFormatting(buf, "", ISO_8859_1, true)
+			n, err = p.WriteFormattedComment(buf, ISO_8859_1)
 		}
 		assert.Equal(t, err, nil)
 		s := string(buf.Bytes())
@@ -942,6 +945,7 @@ func TestWriteCommentWithFormatting(t *testing.T) {
 func TestWriteFormattedCommentWithoutFormatting(t *testing.T) {
 	for _, test := range writeFormattedCommentWithoutFormattingTests {
 		p, err := parse(test.input, true)
+		assert.Equal(t, err, nil)
 
 		buf := new(bytes.Buffer)
 		var n int