Add support for storing comments with keys

lex.go

diff --git a/lex.go b/lex.go
index f6c0c23..1d6b532 100644
--- a/lex.go
+++ b/lex.go
@@ -44,8 +44,9 @@
 const (
 	itemError itemType = iota // error occurred; value is text of error
 	itemEOF
-	itemKey   // a key
-	itemValue // a value
+	itemKey     // a key
+	itemValue   // a value
+	itemComment // a comment
 )
 
 // defines a constant for EOF
@@ -207,6 +208,8 @@
 
 // lexComment scans a comment line. The comment character has already been scanned.
 func lexComment(l *lexer) stateFn {
+	l.acceptRun(whitespace)
+	l.ignore()
 	for {
 		switch r := l.next(); {
 		case isEOF(r):
@@ -214,8 +217,10 @@
 			l.emit(itemEOF)
 			return nil
 		case isEOL(r):
-			l.ignore()
+			l.emit(itemComment)
 			return lexBeforeKey
+		default:
+			l.appendRune(r)
 		}
 	}
 }

parser.go

diff --git a/parser.go b/parser.go
index e6efccc..a7028ca 100644
--- a/parser.go
+++ b/parser.go
@@ -19,22 +19,36 @@
 
 	p.lex = lex(input)
 	properties = NewProperties()
+	key := ""
+	comments := []string{}
 
 	for {
-		token := p.expectOneOf(itemKey, itemEOF)
-		if token.typ == itemEOF {
-			break
+		token := p.expectOneOf(itemComment, itemKey, itemEOF)
+		switch token.typ {
+		case itemEOF:
+			goto done
+		case itemComment:
+			comments = append(comments, token.val)
+			continue
+		case itemKey:
+			key = token.val
 		}
-		key := token.val
 
 		token = p.expectOneOf(itemValue, itemEOF)
-		if token.typ == itemEOF {
-			properties.m[key] = ""
-			break
+		if len(comments) > 0 {
+			properties.c[key] = comments
+			comments = []string{}
 		}
-		properties.m[key] = token.val
+		switch token.typ {
+		case itemEOF:
+			properties.m[key] = ""
+			goto done
+		case itemValue:
+			properties.m[key] = token.val
+		}
 	}
 
+done:
 	return properties, nil
 }
 
@@ -51,12 +65,15 @@
 	return token
 }
 
-func (p *parser) expectOneOf(expected1, expected2 itemType) (token item) {
+func (p *parser) expectOneOf(expected ...itemType) (token item) {
 	token = p.lex.nextItem()
-	if token.typ != expected1 && token.typ != expected2 {
-		p.unexpected(token)
+	for _, v := range expected {
+		if token.typ == v {
+			return token
+		}
 	}
-	return token
+	p.unexpected(token)
+	panic("unexpected token")
 }
 
 func (p *parser) unexpected(token item) {

properties.go

diff --git a/properties.go b/properties.go
index cde68fe..898320e 100644
--- a/properties.go
+++ b/properties.go
@@ -46,6 +46,7 @@
 	Postfix string
 
 	m map[string]string
+	c map[string][]string
 }
 
 // NewProperties creates a new Properties struct with the default
@@ -55,6 +56,7 @@
 		Prefix:  "${",
 		Postfix: "}",
 		m:       make(map[string]string),
+		c:       make(map[string][]string),
 	}
 }
 
@@ -90,6 +92,28 @@
 
 // ----------------------------------------------------------------------------
 
+// GetComment returns the last comment before the given key or an
+// empty string.
+func (p *Properties) GetComment(key string) string {
+	comments, ok := p.c[key]
+	if !ok || len(comments) == 0 {
+		return ""
+	}
+	return comments[len(comments)-1]
+}
+
+// ----------------------------------------------------------------------------
+
+// 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
+	}
+	return nil
+}
+
+// ----------------------------------------------------------------------------
+
 // GetBool checks if the expanded value is one of '1', 'yes',
 // 'true' or 'on' if the key exists. The comparison is case-insensitive.
 // If the key does not exist the default value is returned.

properties_test.go

diff --git a/properties_test.go b/properties_test.go
index fc99dad..f0c6e49 100644
--- a/properties_test.go
+++ b/properties_test.go
@@ -109,8 +109,25 @@
 	{"key=${USER}\nUSER=value", "key", "value", "USER", "value"},
 }
 
-// define error test cases in the form of
-// {"input", "expected error message"}
+var commentTests = []struct {
+	input, key, value string
+	comments          []string
+}{
+	{"key=value", "key", "value", nil},
+	{"#comment\nkey=value", "key", "value", []string{"comment"}},
+	{"# comment\nkey=value", "key", "value", []string{"comment"}},
+	{"#  comment\nkey=value", "key", "value", []string{"comment"}},
+	{"# comment\n\nkey=value", "key", "value", []string{"comment"}},
+	{"# comment1\n# comment2\nkey=value", "key", "value", []string{"comment1", "comment2"}},
+	{"# comment1\n\n# comment2\n\nkey=value", "key", "value", []string{"comment1", "comment2"}},
+	{"!comment\nkey=value", "key", "value", []string{"comment"}},
+	{"! comment\nkey=value", "key", "value", []string{"comment"}},
+	{"!  comment\nkey=value", "key", "value", []string{"comment"}},
+	{"! comment\n\nkey=value", "key", "value", []string{"comment"}},
+	{"! comment1\n! comment2\nkey=value", "key", "value", []string{"comment1", "comment2"}},
+	{"! comment1\n\n! comment2\n\nkey=value", "key", "value", []string{"comment1", "comment2"}},
+}
+
 var errorTests = []struct {
 	input, msg string
 }{
@@ -130,8 +147,6 @@
 	{"key=valu${ke", "Malformed expression"},
 }
 
-// define write encoding test cases in the form of
-// {"input", "expected output after write", ["UTF-8", "ISO-8859-1"]}
 var writeTests = []struct {
 	input, output, encoding string
 }{
@@ -504,6 +519,20 @@
 	c.Assert(func() { p.MustGetString("invalid") }, PanicMatches, "unknown property: invalid")
 }
 
+func (l *TestSuite) TestComment(c *C) {
+	for _, test := range commentTests {
+		p, err := parse(test.input)
+		c.Assert(err, IsNil)
+		c.Assert(p.MustGetString(test.key), Equals, test.value)
+		c.Assert(p.GetComments(test.key), DeepEquals, test.comments)
+		if test.comments != nil {
+			c.Assert(p.GetComment(test.key), Equals, test.comments[len(test.comments)-1])
+		} else {
+			c.Assert(p.GetComment(test.key), Equals, "")
+		}
+	}
+}
+
 func (l *TestSuite) TestFilter(c *C) {
 	for _, test := range filterTests {
 		p, err := parse(test.input)