From 62e2d802edc0351341ea74d9851dc7dd66c3b279 Mon Sep 17 00:00:00 2001
From: Albert Nigmatzianov <albertnigma@gmail.com>
Date: Tue, 21 Mar 2017 14:01:44 +0500
Subject: [PATCH] Fix #141 (#142)

* Use String() of key if it exists during TreeFromMap
---
 tomltree_create.go      | 11 +++++++----
 tomltree_create_test.go | 22 ++++++++++++++++------
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/tomltree_create.go b/tomltree_create.go
index d447755..c9e8b37 100644
--- a/tomltree_create.go
+++ b/tomltree_create.go
@@ -48,6 +48,8 @@ func simpleValueCoercion(object interface{}) (interface{}, error) {
 		return uint64(original), nil
 	case float32:
 		return float64(original), nil
+	case fmt.Stringer:
+		return original.String(), nil
 	default:
 		return nil, fmt.Errorf("cannot convert type %T to TomlTree", object)
 	}
@@ -102,9 +104,10 @@ func toTree(object interface{}) (interface{}, error) {
 		values := map[string]interface{}{}
 		keys := value.MapKeys()
 		for _, key := range keys {
-			k, ok := key.Interface().(string)
-			if !ok {
-				return nil, fmt.Errorf("map key needs to be a string, not %T", key.Interface())
+			if key.Kind() != reflect.String {
+				if _, ok := key.Interface().(string); !ok {
+					return nil, fmt.Errorf("map key needs to be a string, not %T (%v)", key.Interface(), key.Kind())
+				}
 			}
 
 			v := value.MapIndex(key)
@@ -112,7 +115,7 @@ func toTree(object interface{}) (interface{}, error) {
 			if err != nil {
 				return nil, err
 			}
-			values[k] = newValue
+			values[key.String()] = newValue
 		}
 		return &TomlTree{values, Position{}}, nil
 	}
diff --git a/tomltree_create_test.go b/tomltree_create_test.go
index 8523f9e..74bde9e 100644
--- a/tomltree_create_test.go
+++ b/tomltree_create_test.go
@@ -5,6 +5,14 @@ import (
 	"time"
 )
 
+type customString string
+
+type stringer struct{}
+
+func (s stringer) String() string {
+	return "stringer"
+}
+
 func validate(t *testing.T, path string, object interface{}) {
 	switch o := object.(type) {
 	case *TomlTree:
@@ -46,14 +54,16 @@ func TestTomlTreeCreateToTree(t *testing.T) {
 		"uint32":   uint32(2),
 		"float32":  float32(2),
 		"a_bool":   false,
+		"stringer": stringer{},
 		"nested": map[string]interface{}{
 			"foo": "bar",
 		},
-		"array":       []string{"a", "b", "c"},
-		"array_uint":  []uint{uint(1), uint(2)},
-		"array_table": []map[string]interface{}{map[string]interface{}{"sub_map": 52}},
-		"array_times": []time.Time{time.Now(), time.Now()},
-		"map_times":   map[string]time.Time{"now": time.Now()},
+		"array":                 []string{"a", "b", "c"},
+		"array_uint":            []uint{uint(1), uint(2)},
+		"array_table":           []map[string]interface{}{map[string]interface{}{"sub_map": 52}},
+		"array_times":           []time.Time{time.Now(), time.Now()},
+		"map_times":             map[string]time.Time{"now": time.Now()},
+		"custom_string_map_key": map[customString]interface{}{customString("custom"): "custom"},
 	}
 	tree, err := TreeFromMap(data)
 	if err != nil {
@@ -72,7 +82,7 @@ func TestTomlTreeCreateToTreeInvalidLeafType(t *testing.T) {
 
 func TestTomlTreeCreateToTreeInvalidMapKeyType(t *testing.T) {
 	_, err := TreeFromMap(map[string]interface{}{"foo": map[int]interface{}{2: 1}})
-	expected := "map key needs to be a string, not int"
+	expected := "map key needs to be a string, not int (int)"
 	if err.Error() != expected {
 		t.Fatalf("expected error %s, got %s", expected, err.Error())
 	}