From 5fd6e9cce0672a3b4a1a43193c0b194c31096c81 Mon Sep 17 00:00:00 2001
From: Thomas Pelletier <pelletier.thomas@gmail.com>
Date: Sun, 26 Dec 2021 18:29:46 +0100
Subject: [PATCH] Encode: add comment struct tag (#711)

Similar to v1, add a `comment` struct that that makes the encoder emit a comment
before the annotated element, if permitted. Unlike v1, comments are compact by
default (and cannot be changed).

Fixes #595.
---
 marshaler.go      | 49 ++++++++++++++++++++++++++++++++---------------
 marshaler_test.go | 27 ++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 15 deletions(-)

diff --git a/marshaler.go b/marshaler.go
index 8a65941..4aaff21 100644
--- a/marshaler.go
+++ b/marshaler.go
@@ -104,30 +104,31 @@ func (enc *Encoder) SetIndentTables(indent bool) *Encoder {
 // Intermediate tables are always printed.
 //
 // By default, strings are encoded as literal string, unless they contain either
-// a newline character or a single quote. In that case they are emitted as quoted
-// strings.
+// a newline character or a single quote. In that case they are emitted as
+// quoted strings.
 //
 // When encoding structs, fields are encoded in order of definition, with their
 // exact name.
 //
 // Struct tags
 //
-// The encoding of each public struct field can be customized by the
-// format string in the "toml" key of the struct field's tag. This
-// follows encoding/json's convention. The format string starts with
-// the name of the field, optionally followed by a comma-separated
-// list of options. The name may be empty in order to provide options
-// without overriding the default name.
+// The encoding of each public struct field can be customized by the format
+// string in the "toml" key of the struct field's tag. This follows
+// encoding/json's convention. The format string starts with the name of the
+// field, optionally followed by a comma-separated list of options. The name may
+// be empty in order to provide options without overriding the default name.
 //
-// The "multiline" option emits strings as quoted multi-line TOML
-// strings. It has no effect on fields that would not be encoded as
-// strings.
+// The "multiline" option emits strings as quoted multi-line TOML strings. It
+// has no effect on fields that would not be encoded as strings.
 //
-// The "inline" option turns fields that would be emitted as tables
-// into inline tables instead. It has no effect on other fields.
+// The "inline" option turns fields that would be emitted as tables into inline
+// tables instead. It has no effect on other fields.
 //
-// The "omitempty" option prevents empty values or groups from being
-// emitted.
+// The "omitempty" option prevents empty values or groups from being emitted.
+//
+// In addition to the "toml" tag struct tag, a "comment" tag can be used to emit
+// a TOML comment before the value being annotated. Comments are ignored inside
+// inline tables.
 func (enc *Encoder) Encode(v interface{}) error {
 	var (
 		b   []byte
@@ -156,6 +157,7 @@ func (enc *Encoder) Encode(v interface{}) error {
 type valueOptions struct {
 	multiline bool
 	omitempty bool
+	comment   string
 }
 
 type encoderCtx struct {
@@ -306,6 +308,10 @@ func (enc *Encoder) encodeKv(b []byte, ctx encoderCtx, options valueOptions, v r
 		return b, nil
 	}
 
+	if !ctx.inline {
+		b = enc.encodeComment(ctx.indent, options.comment, b)
+	}
+
 	b = enc.indent(ctx.indent, b)
 
 	b, err = enc.encodeKey(b, ctx.key)
@@ -441,6 +447,8 @@ func (enc *Encoder) encodeTableHeader(ctx encoderCtx, b []byte) ([]byte, error)
 		return b, nil
 	}
 
+	b = enc.encodeComment(ctx.indent, ctx.options.comment, b)
+
 	b = enc.indent(ctx.indent, b)
 
 	b = append(b, '[')
@@ -590,6 +598,7 @@ func (enc *Encoder) encodeStruct(b []byte, ctx encoderCtx, v reflect.Value) ([]b
 		options := valueOptions{
 			multiline: opts.multiline,
 			omitempty: opts.omitempty,
+			comment:   fieldType.Tag.Get("comment"),
 		}
 
 		if opts.inline || !willConvertToTableOrArrayTable(ctx, f) {
@@ -602,6 +611,16 @@ func (enc *Encoder) encodeStruct(b []byte, ctx encoderCtx, v reflect.Value) ([]b
 	return enc.encodeTable(b, ctx, t)
 }
 
+func (enc *Encoder) encodeComment(indent int, comment string, b []byte) []byte {
+	if comment != "" {
+		b = enc.indent(indent, b)
+		b = append(b, "# "...)
+		b = append(b, comment...)
+		b = append(b, '\n')
+	}
+	return b
+}
+
 func isValidName(s string) bool {
 	if s == "" {
 		return false
diff --git a/marshaler_test.go b/marshaler_test.go
index 49c2a4f..ae9d1e7 100644
--- a/marshaler_test.go
+++ b/marshaler_test.go
@@ -21,6 +21,12 @@ func TestMarshal(t *testing.T) {
 		A interface{} `toml:",inline"`
 	}
 
+	type comments struct {
+		One   int
+		Two   int   `comment:"Before kv"`
+		Three []int `comment:"Before array"`
+	}
+
 	examples := []struct {
 		desc     string
 		v        interface{}
@@ -535,6 +541,27 @@ J = 42
 K = 42
 L = 2.2`,
 		},
+		{
+			desc: "comments",
+			v: struct {
+				Table comments `comment:"Before table"`
+			}{
+				Table: comments{
+					One:   1,
+					Two:   2,
+					Three: []int{1, 2, 3},
+				},
+			},
+			expected: `
+# Before table
+[Table]
+One = 1
+# Before kv
+Two = 2
+# Before array
+Three = [1, 2, 3]
+`,
+		},
 	}
 
 	for _, e := range examples {