func: Implement TextMarshaler for map keys

Author: thockin

funcr/funcr.go

@@ -36,6 +36,7 @@ package funcr
 
 import (
 	"bytes"
+	"encoding"
 	"fmt"
 	"path/filepath"
 	"reflect"
@@ -492,12 +493,24 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
 			if i > 0 {
 				buf.WriteByte(',')
 			}
-			// prettyWithFlags will produce already-escaped values
-			keystr := f.prettyWithFlags(it.Key().Interface(), 0)
-			if t.Key().Kind() != reflect.String {
-				// JSON only does string keys.  Unlike Go's standard JSON, we'll
-				// convert just about anything to a string.
+			// If a map key supports TextMarshaler, use it.
+			keystr := ""
+			if m, ok := it.Key().Interface().(encoding.TextMarshaler); ok {
+				txt, err := m.MarshalText()
+				if err != nil {
+					keystr = fmt.Sprintf("<error-MarshalText: %s>", err.Error())
+				} else {
+					keystr = string(txt)
+				}
 				keystr = prettyString(keystr)
+			} else {
+				// prettyWithFlags will produce already-escaped values
+				keystr = f.prettyWithFlags(it.Key().Interface(), 0)
+				if t.Key().Kind() != reflect.String {
+					// JSON only does string keys.  Unlike Go's standard JSON, we'll
+					// convert just about anything to a string.
+					keystr = prettyString(keystr)
+				}
 			}
 			buf.WriteString(keystr)
 			buf.WriteByte(':')

funcr/funcr_test.go

@@ -37,6 +37,20 @@ func ptrstr(s string) *string {
 	return &s
 }
 
+// point implements encoding.TextMarshaler and can be used as a map key.
+type point struct{ x, y int }
+
+func (p point) MarshalText() ([]byte, error) {
+	return []byte(fmt.Sprintf("(%d, %d)", p.x, p.y)), nil
+}
+
+// pointErr implements encoding.TextMarshaler but returns an error.
+type pointErr struct{ x, y int }
+
+func (p pointErr) MarshalText() ([]byte, error) {
+	return nil, fmt.Errorf("uh oh: %d, %d", p.x, p.y)
+}
+
 // Logging this should result in the MarshalLog() value.
 type Tmarshaler string
 
@@ -288,6 +302,17 @@ func TestPretty(t *testing.T) {
 			},
 			exp: `{"9.5":3}`,
 		},
+		{
+			val: map[point]int{
+				{x: 1, y: 2}: 3,
+			},
+		},
+		{
+			val: map[pointErr]int{
+				{x: 1, y: 2}: 3,
+			},
+			exp: `{"<error-MarshalText: uh oh: 1, 2>":3}`,
+		},
 		{
 			val: struct {
 				X int `json:"x"`
@@ -496,7 +521,7 @@ func TestPretty(t *testing.T) {
 		} else {
 			jb, err := json.Marshal(tc.val)
 			if err != nil {
-				t.Fatalf("[%d]: unexpected error: %v", i, err)
+				t.Fatalf("[%d]: unexpected error: %v\ngot: %q", i, err, ours)
 			}
 			want = string(jb)
 		}