add Marshaler interface

encode.go

@@ -38,6 +38,12 @@ var quotedReplacer = strings.NewReplacer(
 	"\\", "\\\\",
 )
 
+// Marshaler is the interface implemented by objects that can marshal
+// themselves to a TOML representation.
+type Marshaler interface {
+	MarshalTOML() ([]byte, error)
+}
+
 // Encoder controls the encoding of Go values to a TOML document to some
 // io.Writer.
 //
@@ -114,6 +120,10 @@ func (enc *Encoder) encode(key Key, rv reflect.Value) {
 	case time.Time, TextMarshaler:
 		enc.keyEqElement(key, rv)
 		return
+
+	case Marshaler:
+		enc.eMarshaler(key, rv, enc.encode)
+		return
 	}
 
 	k := rv.Kind()
@@ -253,6 +263,11 @@ func (enc *Encoder) eTable(key Key, rv reflect.Value) {
 }
 
 func (enc *Encoder) eMapOrStruct(key Key, rv reflect.Value) {
+	if implementsMarshaler(rv) {
+		enc.eMarshaler(key, rv, enc.eMap)
+		return
+	}
+
 	switch rv := eindirect(rv); rv.Kind() {
 	case reflect.Map:
 		enc.eMap(key, rv)
@@ -537,6 +552,30 @@ func encPanic(err error) {
 	panic(tomlEncodeError{err})
 }
 
+func (enc *Encoder) eMarshaler(
+	key Key, rv reflect.Value, next func(Key, reflect.Value),
+) {
+	buf, err := rv.Interface().(Marshaler).MarshalTOML()
+	if err != nil {
+		encPanic(err)
+		return
+	}
+
+	v := map[string]interface{}{}
+	_, err = Decode(string(buf), &v)
+	if err != nil {
+		encPanic(err)
+		return
+	}
+
+	next(key, eindirect(reflect.ValueOf(v)))
+}
+
+func implementsMarshaler(rv reflect.Value) bool {
+	_, ok := rv.Interface().(Marshaler)
+	return ok
+}
+
 func eindirect(v reflect.Value) reflect.Value {
 	switch v.Kind() {
 	case reflect.Ptr, reflect.Interface:

encode_test.go

@@ -531,6 +531,83 @@ func TestEncodeAnonymousStruct(t *testing.T) {
 	encodeExpected(t, "embedded anonymous tagged struct", v1, expected, nil)
 }
 
+type withMarshalTOML struct{}
+
+func (withMarshalTOML) MarshalTOML() ([]byte, error) {
+	return []byte("x = 1\ny = 1\n"), nil
+}
+
+func TestEncodeMarshalerAnonymousStruct(t *testing.T) {
+	type Inner struct {
+		N int
+		withMarshalTOML
+	}
+	type Outer0 struct{ Inner }
+	type Outer1 struct {
+		Inner `toml:"inner"`
+		X     int
+	}
+
+	v0 := Inner{}
+	expected := "x = 1\ny = 1\n"
+	encodeExpected(t, "embedded struct with MarshalTOML", v0, expected, nil)
+
+	v1 := Outer0{Inner{}}
+	expected = "x = 1\ny = 1\n"
+	encodeExpected(t, "embedded anonymous untagged struct with MarshalTOML", v1, expected, nil)
+
+	v2 := Outer1{Inner{}, 3}
+	expected = "x = 1\ny = 1\n"
+	encodeExpected(t, "embedded anonymous tagged struct with MarshalTOML", v2, expected, nil)
+}
+
+func TestEncodeMarshalerField(t *testing.T) {
+	type Inner struct {
+		withMarshalTOML
+		N int
+		Z int
+	}
+	type Outer struct {
+		A     int
+		Field Inner
+	}
+	type Outer1 struct {
+		B   int
+		Ins []Inner `toml:"slice"`
+	}
+
+	v0 := Outer{
+		A:     1,
+		Field: Inner{},
+	}
+	expected := `A = 1
+
+[Field]
+  x = 1
+  y = 1
+`
+	encodeExpected(t, "struct field with MarshalTOML", v0, expected, nil)
+
+	v1 := Outer1{
+		B: 1,
+		Ins: []Inner{
+			Inner{},
+			Inner{},
+		},
+	}
+	expected = `B = 1
+
+[[slice]]
+  x = 1
+  y = 1
+
+[[slice]]
+  x = 1
+  y = 1
+`
+	encodeExpected(t, "struct slice field with MarshalTOML", v1, expected, nil)
+}
+
 func TestEncodeAnonymousStructPointerField(t *testing.T) {
 	type Inner struct{ N int }
 	type Outer0 struct{ *Inner }