瀏覽代碼

Implement SetDefaults for maps. Skip proto3 scalar fields instead of panicking.

David Symonds 11 年之前
父節點
當前提交
de8c523226
共有 4 個文件被更改,包括 93 次插入13 次删除
  1. 11 0
      proto/all_test.go
  2. 44 13
      proto/lib.go
  3. 6 0
      proto/proto3_proto/proto3.proto
  4. 32 0
      proto/proto3_test.go

+ 11 - 0
proto/all_test.go

@@ -1442,6 +1442,17 @@ func TestSetDefaultsWithRepeatedSubMessage(t *testing.T) {
 	}
 }
 
+func TestSetDefaultWithRepeatedNonMessage(t *testing.T) {
+	m := &MyMessage{
+		Pet: []string{"turtle", "wombat"},
+	}
+	expected := Clone(m)
+	SetDefaults(m)
+	if !Equal(m, expected) {
+		t.Errorf("\n got %v\nwant %v", m, expected)
+	}
+}
+
 func TestMaximumTagNumber(t *testing.T) {
 	m := &MaxTag{
 		LastField: String("natural goat essence"),

+ 44 - 13
proto/lib.go

@@ -607,13 +607,15 @@ func setDefaults(v reflect.Value, recur, zeros bool) {
 
 	for _, ni := range dm.nested {
 		f := v.Field(ni)
-		if f.IsNil() {
-			continue
-		}
-		// f is *T or []*T
-		if f.Kind() == reflect.Ptr {
+		// f is *T or []*T or map[T]*T
+		switch f.Kind() {
+		case reflect.Ptr:
+			if f.IsNil() {
+				continue
+			}
 			setDefaults(f, recur, zeros)
-		} else {
+
+		case reflect.Slice:
 			for i := 0; i < f.Len(); i++ {
 				e := f.Index(i)
 				if e.IsNil() {
@@ -621,6 +623,15 @@ func setDefaults(v reflect.Value, recur, zeros bool) {
 				}
 				setDefaults(e, recur, zeros)
 			}
+
+		case reflect.Map:
+			for _, k := range f.MapKeys() {
+				e := f.MapIndex(k)
+				if e.IsNil() {
+					continue
+				}
+				setDefaults(e, recur, zeros)
+			}
 		}
 	}
 }
@@ -646,10 +657,6 @@ type scalarField struct {
 	value interface{}  // the proto-declared default value, or nil
 }
 
-func ptrToStruct(t reflect.Type) bool {
-	return t.Kind() == reflect.Ptr && t.Elem().Kind() == reflect.Struct
-}
-
 // t is a struct type.
 func buildDefaultMessage(t reflect.Type) (dm defaultMessage) {
 	sprop := GetProperties(t)
@@ -661,9 +668,33 @@ func buildDefaultMessage(t reflect.Type) (dm defaultMessage) {
 		}
 		ft := t.Field(fi).Type
 
-		// nested messages
-		if ptrToStruct(ft) || (ft.Kind() == reflect.Slice && ptrToStruct(ft.Elem())) {
-			dm.nested = append(dm.nested, fi)
+		var canHaveDefault, nestedMessage bool
+		switch ft.Kind() {
+		case reflect.Ptr:
+			if ft.Elem().Kind() == reflect.Struct {
+				nestedMessage = true
+			} else {
+				canHaveDefault = true // proto2 scalar field
+			}
+
+		case reflect.Slice:
+			switch ft.Elem().Kind() {
+			case reflect.Ptr:
+				nestedMessage = true // repeated message
+			case reflect.Uint8:
+				canHaveDefault = true // bytes field
+			}
+
+		case reflect.Map:
+			if ft.Elem().Kind() == reflect.Ptr {
+				nestedMessage = true // map with message values
+			}
+		}
+
+		if !canHaveDefault {
+			if nestedMessage {
+				dm.nested = append(dm.nested, fi)
+			}
 			continue
 		}
 

+ 6 - 0
proto/proto3_proto/proto3.proto

@@ -31,6 +31,8 @@
 
 syntax = "proto3";
 
+import "testdata/test.proto";
+
 package proto3_proto;
 
 message Message {
@@ -51,6 +53,10 @@ message Message {
 
   repeated uint64 key = 5;
   Nested nested = 6;
+
+  map<string, Nested> terrain = 10;
+  testdata.SubDefaults proto2_field = 11;
+  map<string, testdata.SubDefaults> proto2_value = 13;
 }
 
 message Nested {

+ 32 - 0
proto/proto3_test.go

@@ -36,6 +36,7 @@ import (
 
 	"github.com/golang/protobuf/proto"
 	pb "github.com/golang/protobuf/proto/proto3_proto"
+	tpb "github.com/golang/protobuf/proto/testdata"
 )
 
 func TestProto3ZeroValues(t *testing.T) {
@@ -91,3 +92,34 @@ func TestRoundTripProto3(t *testing.T) {
 		t.Errorf("proto.Equal returned false:\n m: %v\nm2: %v", m, m2)
 	}
 }
+
+func TestProto3SetDefaults(t *testing.T) {
+	in := &pb.Message{
+		Terrain: map[string]*pb.Nested{
+			"meadow": new(pb.Nested),
+		},
+		Proto2Field: new(tpb.SubDefaults),
+		Proto2Value: map[string]*tpb.SubDefaults{
+			"badlands": new(tpb.SubDefaults),
+		},
+	}
+
+	got := proto.Clone(in).(*pb.Message)
+	proto.SetDefaults(got)
+
+	// There are no defaults in proto3.  Everything should be the zero value, but
+	// we need to remember to set defaults for nested proto2 messages.
+	want := &pb.Message{
+		Terrain: map[string]*pb.Nested{
+			"meadow": new(pb.Nested),
+		},
+		Proto2Field: &tpb.SubDefaults{N: proto.Int64(7)},
+		Proto2Value: map[string]*tpb.SubDefaults{
+			"badlands": &tpb.SubDefaults{N: proto.Int64(7)},
+		},
+	}
+
+	if !proto.Equal(got, want) {
+		t.Errorf("with in = %v\nproto.SetDefaults(in) =>\ngot %v\nwant %v", in, got, want)
+	}
+}