Merge pull request #21 from hermannm/textmarshaler-colorization

Implement colorization for encoding.TextMarshaler

encode.go

@@ -911,7 +911,7 @@ func (e encoder) encodeTextMarshaler(b []byte, p unsafe.Pointer, t reflect.Type,
 	switch v.Kind() {
 	case reflect.Ptr, reflect.Interface:
 		if v.IsNil() {
-			return append(b, `null`...), nil
+			return e.clrs.appendNull(b), nil
 		}
 	}
 
@@ -920,7 +920,14 @@ func (e encoder) encodeTextMarshaler(b []byte, p unsafe.Pointer, t reflect.Type,
 		return b, err
 	}
 
-	return e.doEncodeString(b, unsafe.Pointer(&s))
+	if e.clrs == nil {
+		return e.doEncodeString(b, unsafe.Pointer(&s))
+	}
+
+	b = append(b, e.clrs.TextMarshaler...)
+	b, err = e.doEncodeString(b, unsafe.Pointer(&s))
+	b = append(b, ansiReset...)
+	return b, err
 }
 
 func appendCompactEscapeHTML(dst []byte, src []byte) []byte {

helper/fatihcolor/fatihcolor.go

@@ -37,19 +37,23 @@ type Colors struct {
 	// Punc is the color for punctuation such as colons, braces, etc.
 	// Frequently Punc will just be color.Bold.
 	Punc *color.Color
+
+	// TextMarshaler is the color for types implementing encoding.TextMarshaler.
+	TextMarshaler *color.Color
 }
 
 // DefaultColors returns default Colors instance.
 func DefaultColors() *Colors {
 	return &Colors{
-		Bool:     color.New(color.FgYellow),
-		Bytes:    color.New(color.Faint),
-		Datetime: color.New(color.FgGreen, color.Faint),
-		Key:      color.New(color.FgBlue, color.Bold),
-		Null:     color.New(color.Faint),
-		Number:   color.New(color.FgCyan),
-		String:   color.New(color.FgGreen),
-		Punc:     color.New(color.Bold),
+		Bool:          color.New(color.FgYellow),
+		Bytes:         color.New(color.Faint),
+		Datetime:      color.New(color.FgGreen, color.Faint),
+		Key:           color.New(color.FgBlue, color.Bold),
+		Null:          color.New(color.Faint),
+		Number:        color.New(color.FgCyan),
+		String:        color.New(color.FgGreen),
+		Punc:          color.New(color.Bold),
+		TextMarshaler: color.New(color.FgGreen),
 	}
 }
 
@@ -60,14 +64,15 @@ func ToCoreColors(clrs *Colors) *jsoncolor.Colors {
 	}
 
 	return &jsoncolor.Colors{
-		Null:   ToCoreColor(clrs.Null),
-		Bool:   ToCoreColor(clrs.Bool),
-		Number: ToCoreColor(clrs.Number),
-		String: ToCoreColor(clrs.String),
-		Key:    ToCoreColor(clrs.Key),
-		Bytes:  ToCoreColor(clrs.Bytes),
-		Time:   ToCoreColor(clrs.Datetime),
-		Punc:   ToCoreColor(clrs.Punc),
+		Null:          ToCoreColor(clrs.Null),
+		Bool:          ToCoreColor(clrs.Bool),
+		Number:        ToCoreColor(clrs.Number),
+		String:        ToCoreColor(clrs.String),
+		Key:           ToCoreColor(clrs.Key),
+		Bytes:         ToCoreColor(clrs.Bytes),
+		Time:          ToCoreColor(clrs.Datetime),
+		Punc:          ToCoreColor(clrs.Punc),
+		TextMarshaler: ToCoreColor(clrs.TextMarshaler),
 	}
 }
 

jsoncolor.go

@@ -36,6 +36,9 @@ type Colors struct {
 
 	// Punc is the color for JSON punctuation: []{},: etc.
 	Punc Color
+
+	// TextMarshaler is the color for values implementing encoding.TextMarshaler.
+	TextMarshaler Color
 }
 
 // appendNull appends a colorized "null" to b.
@@ -131,14 +134,15 @@ const ansiReset = "\x1b[0m"
 // with some deviation.
 func DefaultColors() *Colors {
 	return &Colors{
-		Null:   Color("\x1b[2m"),
-		Bool:   Color("\x1b[1m"),
-		Number: Color("\x1b[36m"),
-		String: Color("\x1b[32m"),
-		Key:    Color("\x1b[34;1m"),
-		Bytes:  Color("\x1b[2m"),
-		Time:   Color("\x1b[32;2m"),
-		Punc:   Color{}, // No colorization
+		Null:          Color("\x1b[2m"),
+		Bool:          Color("\x1b[1m"),
+		Number:        Color("\x1b[36m"),
+		String:        Color("\x1b[32m"),
+		Key:           Color("\x1b[34;1m"),
+		Bytes:         Color("\x1b[2m"),
+		Time:          Color("\x1b[32;2m"),
+		Punc:          Color{},           // No colorization
+		TextMarshaler: Color("\x1b[32m"), // Same as String
 	}
 }
 

jsoncolor_test.go

@@ -573,3 +573,26 @@ func TestEquivalenceRecords(t *testing.T) {
 	err = jsoncolor.NewEncoder(bufJ).Encode(rec)
 	require.Equal(t, bufStdj.String(), bufJ.String())
 }
+
+// TextMarshaler implements encoding.TextMarshaler
+type TextMarshaler struct {
+	Text string
+}
+
+func (t TextMarshaler) MarshalText() ([]byte, error) {
+	return []byte(t.Text), nil
+}
+
+func TestEncode_TextMarshaler(t *testing.T) {
+	buf := &bytes.Buffer{}
+	enc := jsoncolor.NewEncoder(buf)
+	enc.SetColors(&jsoncolor.Colors{
+		TextMarshaler: jsoncolor.Color("\x1b[36m"),
+	})
+
+	text := TextMarshaler{Text: "example text"}
+
+	require.NoError(t, enc.Encode(text))
+	require.Equal(t, "\x1b[36m\"example text\"\x1b[0m\n", buf.String(),
+		"expected TextMarshaler encoding to use Colors.TextMarshaler")
+}