From 9111f3b43928628990d3e28d627b4b6fa23f0435 Mon Sep 17 00:00:00 2001
From: Herbie Ong <herbie@google.com>
Date: Fri, 6 Sep 2019 14:35:09 -0700
Subject: [PATCH] encoding/protojson: add MarshalOptions.UseEnumNumbers

UseEnumNumbers=true will emit enum values as JSON numbers.

Change-Id: I6f3c814e06dc1e3dd595ad35aa79871a49718cd5
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/194017
Reviewed-by: Joe Tsai <thebrokentoaster@gmail.com>
---
 encoding/protojson/encode.go      | 19 +++++++-----
 encoding/protojson/encode_test.go | 50 +++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 7 deletions(-)

diff --git a/encoding/protojson/encode.go b/encoding/protojson/encode.go
index c761fec5..d55786e1 100644
--- a/encoding/protojson/encode.go
+++ b/encoding/protojson/encode.go
@@ -34,6 +34,9 @@ type MarshalOptions struct {
 	// Marshal will return error if there are any missing required fields.
 	AllowPartial bool
 
+	// UseEnumNumbers emits enum values as numbers.
+	UseEnumNumbers bool
+
 	// EmitUnpopulated specifies whether to emit unpopulated fields. It does not
 	// emit unpopulated oneof fields or unpopulated extension fields.
 	// The JSON value emitted for unpopulated fields are as follows:
@@ -197,14 +200,16 @@ func (o MarshalOptions) marshalSingular(val pref.Value, fd pref.FieldDescriptor)
 	case pref.EnumKind:
 		if fd.Enum().FullName() == "google.protobuf.NullValue" {
 			o.encoder.WriteNull()
-		} else if desc := fd.Enum().Values().ByNumber(val.Enum()); desc != nil {
-			err := o.encoder.WriteString(string(desc.Name()))
-			if err != nil {
-				return err
-			}
 		} else {
-			// Use numeric value if there is no enum value descriptor.
-			o.encoder.WriteInt(int64(val.Enum()))
+			desc := fd.Enum().Values().ByNumber(val.Enum())
+			if o.UseEnumNumbers || desc == nil {
+				o.encoder.WriteInt(int64(val.Enum()))
+			} else {
+				err := o.encoder.WriteString(string(desc.Name()))
+				if err != nil {
+					return err
+				}
+			}
 		}
 
 	case pref.MessageKind, pref.GroupKind:
diff --git a/encoding/protojson/encode_test.go b/encoding/protojson/encode_test.go
index 661ee59e..4ee275a5 100644
--- a/encoding/protojson/encode_test.go
+++ b/encoding/protojson/encode_test.go
@@ -2124,6 +2124,56 @@ func TestMarshal(t *testing.T) {
       "optNested": null
     }
   ]
+}`,
+	}, {
+		desc: "UseEnumNumbers in singular field",
+		mo:   protojson.MarshalOptions{UseEnumNumbers: true},
+		input: &pb2.Enums{
+			OptEnum:       pb2.Enum_ONE.Enum(),
+			OptNestedEnum: pb2.Enums_UNO.Enum(),
+		},
+		want: `{
+  "optEnum": 1,
+  "optNestedEnum": 1
+}`,
+	}, {
+		desc: "UseEnumNumbers in repeated field",
+		mo:   protojson.MarshalOptions{UseEnumNumbers: true},
+		input: &pb2.Enums{
+			RptEnum:       []pb2.Enum{pb2.Enum_ONE, 2, pb2.Enum_TEN, 42},
+			RptNestedEnum: []pb2.Enums_NestedEnum{pb2.Enums_UNO, pb2.Enums_DOS, 47},
+		},
+		want: `{
+  "rptEnum": [
+    1,
+    2,
+    10,
+    42
+  ],
+  "rptNestedEnum": [
+    1,
+    2,
+    47
+  ]
+}`,
+	}, {
+		desc: "UseEnumNumbers in map field",
+		mo:   protojson.MarshalOptions{UseEnumNumbers: true},
+		input: &pb3.Maps{
+			Uint64ToEnum: map[uint64]pb3.Enum{
+				1:  pb3.Enum_ONE,
+				2:  pb3.Enum_TWO,
+				10: pb3.Enum_TEN,
+				47: 47,
+			},
+		},
+		want: `{
+  "uint64ToEnum": {
+    "1": 1,
+    "2": 2,
+    "10": 10,
+    "47": 47
+  }
 }`,
 	}}