Add preserveFormatting option for comments/whitespace

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)
 	}

lex.go

@@ -60,6 +60,7 @@ type stateFn func(*lexer) stateFn
 
 // lexer holds the state of the scanner.
 type lexer struct {
+	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,13 +164,11 @@ func (l *lexer) nextItem() item {
 
 // lex creates a new scanner for the input string.
 func lex(input string) *lexer {
-	l := &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.
@@ -188,13 +187,29 @@ 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):
 		l.ignore()
 		return lexBeforeKey
 
 	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 isWhitespace(r) && l.keepWS:
+		// treat as part of the next key's comments block
+		// add the whitespace rune r as part of the comment's prefix
+		l.appendRune(r)
+		return lexBeforeKey
+
 	case isWhitespace(r):
 		l.ignore()
 		return lexBeforeKey
@@ -207,12 +222,16 @@ 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.keepWS {
+		l.acceptRun(whitespace)
+		l.ignore()
+	}
 	for {
 		switch r := l.next(); {
 		case isEOF(r):
-			l.ignore()
+			if ! l.keepWS {
+				l.ignore()
+			}
 			l.emit(itemEOF)
 			return nil
 		case isEOL(r):

load.go

@@ -45,6 +45,11 @@ 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.
@@ -146,7 +151,7 @@ func (l *Loader) LoadURL(url string) (*Properties, error) {
 }
 
 func (l *Loader) loadBytes(buf []byte, enc Encoding) (*Properties, error) {
-	p, err := parse(convert(buf, enc))
+	p, err := parse(convert(buf, enc), l.PreserveFormatting)
 	if err != nil {
 		return nil, err
 	}

parser.go

@@ -7,30 +7,60 @@ 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) {
+	l := lex(input)
+	l.keepWS = preserveFormatting
+	go l.run()
+	p := &parser{lex: l}
 	defer p.recover(&err)
 
 	properties = NewProperties()
 	key := ""
-	comments := []string{}
+	comments := []prefixedComment{}
 
 	for {
 		token := p.expectOneOf(itemComment, itemKey, itemEOF)
 		switch token.typ {
 		case itemEOF:
+			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:
-			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 := prefixedComment{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 +69,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 = []prefixedComment{}
 		}
 		switch token.typ {
 		case itemEOF:

properties.go

@@ -48,6 +48,11 @@ func PanicHandler(err error) {
 
 // -----------------------------------------------------------------------------
 
+type prefixedComment struct {
+	prefix string // prefix type (#, !)
+	val string    // The comment string after the prefix
+}
+
 // 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][]prefixedComment
 
 	// Stores the keys in order of appearance.
 	k []string
+
+	// Stores any trailing comments if we want to preserveFormatting
+	trailingComments []prefixedComment
 }
 
 // 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][]prefixedComment{},
 		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][]prefixedComment{}
 }
 
 // ----------------------------------------------------------------------------
@@ -142,15 +150,19 @@ func (p *Properties) GetComment(key string) string {
 	if !ok || len(comments) == 0 {
 		return ""
 	}
-	return comments[len(comments)-1]
+	return comments[len(comments)-1].val
 }
 
 // ----------------------------------------------------------------------------
 
 // 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,19 +171,38 @@ 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] = []prefixedComment{{"#", comment}}
+}
+
+// ----------------------------------------------------------------------------
+
+// SetComment sets the comment for the key.
+func (p *Properties) SetCommentWithPrefix(key, prefix string, comment string) {
+	p.c[key] = []prefixedComment{{prefix, comment}}
 }
 
 // ----------------------------------------------------------------------------
 
 // 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([]prefixedComment, 0, len(comments))
+	for _, comment := range comments {
+		list = append(list, prefixedComment{prefix, comment})
+	}
+	p.c[key] = list
 }
 
 // ----------------------------------------------------------------------------
@@ -592,7 +623,7 @@ 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
 // the properties file format. Otherwise, the properties parser will not be
@@ -609,7 +640,7 @@ func (p *Properties) WriteComment(w io.Writer, prefix string, enc Encoding) (n i
 				// 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
 					}
@@ -626,11 +657,14 @@ func (p *Properties) WriteComment(w io.Writer, prefix string, enc Encoding) (n i
 					}
 
 					for _, c := range comments {
-						x, err = fmt.Fprintf(w, "%s%s\n", prefix, encode(c, "", enc))
-						if err != nil {
-							return
+						comment := strings.TrimSpace(c.val)
+						if comment != "" {
+							x, err = fmt.Fprintf(w, "%s%s\n", prefix, encode(comment, "", enc))
+							if err != nil {
+								return
+							}
+							n += x
 						}
-						n += x
 					}
 				}
 			}
@@ -645,6 +679,48 @@ func (p *Properties) WriteComment(w io.Writer, prefix string, enc Encoding) (n i
 	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
+			}
+		}
+
+		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
+}
+
 // Map returns a copy of the properties as a map.
 func (p *Properties) Map() map[string]string {
 	m := make(map[string]string)