goprotobuf: Support encoding.TextMarshaler and encoding.TextUnmarshaler.

LGTM=r
R=r
CC=golang-codereviews
https://codereview.appspot.com/65740044

4 files changed, 84 insertions, 2 deletions

diff --git a/proto/text.go b/proto/text.go
index f8cb9c9..b4a5dbb 100644
--- a/proto/text.go
+++ b/proto/text.go
@@ -74,6 +74,13 @@ type textWriter struct {
 	w        writer
 }
 
+// textMarshaler is implemented by Messages that can marshal themsleves.
+// It is identical to encoding.TextMarshaler, introduced in go 1.2,
+// which will eventually replace it.
+type textMarshaler interface {
+	MarshalText() (text []byte, err error)
+}
+
 func (w *textWriter) WriteString(s string) (n int, err error) {
 	if !strings.Contains(s, "\n") {
 		if !w.compact && w.complete {
@@ -342,7 +349,15 @@ func writeAny(w *textWriter, v reflect.Value, props *Properties) error {
 			}
 		}
 		w.indent()
-		if err := writeStruct(w, v); err != nil {
+		if tm, ok := v.Interface().(textMarshaler); ok {
+			text, err := tm.MarshalText()
+			if err != nil {
+				return err
+			}
+			if _, err = w.Write(text); err != nil {
+				return err
+			}
+		} else if err := writeStruct(w, v); err != nil {
 			return err
 		}
 		w.unindent()
@@ -629,6 +644,19 @@ func marshalText(w io.Writer, pb Message, compact bool) error {
 		compact:  compact,
 	}
 
+	if tm, ok := pb.(textMarshaler); ok {
+		text, err := tm.MarshalText()
+		if err != nil {
+			return err
+		}
+		if _, err = aw.Write(text); err != nil {
+			return err
+		}
+		if bw != nil {
+			return bw.Flush()
+		}
+		return nil
+	}
 	// Dereference the received pointer so we don't have outer < and >.
 	v := reflect.Indirect(val)
 	if err := writeStruct(aw, v); err != nil {
@@ -642,7 +670,9 @@ func marshalText(w io.Writer, pb Message, compact bool) error {
 
 // MarshalText writes a given protocol buffer in text format.
 // The only errors returned are from w.
-func MarshalText(w io.Writer, pb Message) error { return marshalText(w, pb, false) }
+func MarshalText(w io.Writer, pb Message) error {
+	return marshalText(w, pb, false)
+}
 
 // MarshalTextString is the same as MarshalText, but returns the string directly.
 func MarshalTextString(pb Message) string {
diff --git a/proto/text_parser.go b/proto/text_parser.go
index 13827f6..4d4167c 100644
--- a/proto/text_parser.go
+++ b/proto/text_parser.go
@@ -43,6 +43,13 @@ import (
 	"unicode/utf8"
 )
 
+// textUnmarshaler is implemented by Messages that can unmarshal themsleves.
+// It is identical to encoding.TextUnmarshaler, introduced in go 1.2,
+// which will eventually replace it.
+type textUnmarshaler interface {
+	UnmarshalText(text []byte) error
+}
+
 type ParseError struct {
 	Message string
 	Line    int // 1-based line number
@@ -643,6 +650,7 @@ func (p *textParser) readAny(v reflect.Value, props *Properties) *ParseError {
 		default:
 			return p.errorf("expected '{' or '<', found %q", tok.value)
 		}
+		// TODO: Handle nested messages which implement textUnmarshaler.
 		return p.readStruct(fv, terminator)
 	case reflect.Uint32:
 		if x, err := strconv.ParseUint(tok.value, 0, 32); err == nil {
@@ -661,6 +669,10 @@ func (p *textParser) readAny(v reflect.Value, props *Properties) *ParseError {
 // UnmarshalText reads a protocol buffer in Text format. UnmarshalText resets pb
 // before starting to unmarshal, so any existing data in pb is always removed.
 func UnmarshalText(s string, pb Message) error {
+	if um, ok := pb.(textUnmarshaler); ok {
+		err := um.UnmarshalText([]byte(s))
+		return err
+	}
 	pb.Reset()
 	v := reflect.ValueOf(pb)
 	if pe := newTextParser(s).readStruct(v.Elem(), ""); pe != nil {
diff --git a/proto/text_parser_test.go b/proto/text_parser_test.go
index 5949065..e447ffa 100644
--- a/proto/text_parser_test.go
+++ b/proto/text_parser_test.go
@@ -413,6 +413,16 @@ func TestUnmarshalText(t *testing.T) {
 	}
 }
 
+func TestUnmarshalTextCustomMessage(t *testing.T) {
+	msg := &textMessage{}
+	if err := UnmarshalText("custom", msg); err != nil {
+		t.Errorf("Unexpected error from custom unmarshal: %v", err)
+	}
+	if UnmarshalText("not custom", msg) == nil {
+		t.Errorf("Didn't get expected error from custom unmarshal")
+	}
+}
+
 // Regression test; this caused a panic.
 func TestRepeatedEnum(t *testing.T) {
 	pb := new(RepeatedEnum)
diff --git a/proto/text_test.go b/proto/text_test.go
index c64b073..4b76720 100644
--- a/proto/text_test.go
+++ b/proto/text_test.go
@@ -44,6 +44,26 @@ import (
 	pb "./testdata"
 )
 
+// textMessage implements the methods that allow it to marshal and unmarshal
+// itself as text.
+type textMessage struct {
+}
+
+func (*textMessage) MarshalText() ([]byte, error) {
+	return []byte("custom"), nil
+}
+
+func (*textMessage) UnmarshalText(bytes []byte) error {
+	if string(bytes) != "custom" {
+		return errors.New("expected 'custom'")
+	}
+	return nil
+}
+
+func (*textMessage) Reset()         {}
+func (*textMessage) String() string { return "" }
+func (*textMessage) ProtoMessage()  {}
+
 func newTestMessage() *pb.MyMessage {
 	msg := &pb.MyMessage{
 		Count: proto.Int32(42),
@@ -153,6 +173,16 @@ func TestMarshalText(t *testing.T) {
 	}
 }
 
+func TestMarshalTextCustomMessage(t *testing.T) {
+	buf := new(bytes.Buffer)
+	if err := proto.MarshalText(buf, &textMessage{}); err != nil {
+		t.Fatalf("proto.MarshalText: %v", err)
+	}
+	s := buf.String()
+	if s != "custom" {
+		t.Errorf("Got %q, expected %q", s, "custom")
+	}
+}
 func TestMarshalTextNil(t *testing.T) {
 	want := "<nil>"
 	tests := []proto.Message{nil, (*pb.MyMessage)(nil)}