descriptor: add Enum interface and ForEnum function

descriptor/descriptor.go

@@ -41,6 +41,7 @@ import (
 	"compress/gzip"
 	"fmt"
 	"io/ioutil"
+	"reflect"
 
 	"github.com/golang/protobuf/proto"
 	protobuf "github.com/golang/protobuf/protoc-gen-go/descriptor"
@@ -91,3 +92,37 @@ func ForMessage(msg Message) (fd *protobuf.FileDescriptorProto, md *protobuf.Des
 	}
 	return fd, md
 }
+
+// Enum declares a method to return an protobuf enum's descriptor.
+//
+// Enum types generated by the protocol compiler always satisfy
+// the Enum interface.
+type Enum interface {
+	EnumDescriptor() ([]byte, []int)
+	String() string
+}
+
+// ForEnum returns a FileDescriptorProto and an EnumDescriptorProto from within
+// it describing the given enum.
+func ForEnum(e Enum) (*protobuf.FileDescriptorProto, *protobuf.EnumDescriptorProto) {
+	// EnumDescriptor isn't defined for pointer receivers.
+	if re := reflect.ValueOf(e); re.Kind() == reflect.Ptr && re.IsNil() {
+		e = reflect.Zero(re.Type().Elem()).Interface().(Enum)
+	}
+
+	gz, path := e.EnumDescriptor()
+	fd, err := extractFile(gz)
+	if err != nil {
+		panic(fmt.Sprintf("invalid FileDescriptorProto for %T: %v", e, err))
+	}
+
+	if len(path) == 1 {
+		return fd, fd.EnumType[path[0]]
+	}
+
+	md := fd.MessageType[path[0]]
+	for _, i := range path[1 : len(path)-1] {
+		md = md.NestedType[i]
+	}
+	return fd, md.EnumType[path[len(path)-1]]
+}

descriptor/descriptor_test.go

@@ -20,6 +20,20 @@ func TestMessage(t *testing.T) {
 	}
 }
 
+func TestEnum(t *testing.T) {
+	var enum *protobuf.FieldDescriptorProto_Type
+	fd, ed := descriptor.ForEnum(enum)
+	if pkg, want := fd.GetPackage(), "google.protobuf"; pkg != want {
+		t.Errorf("descriptor.ForEnum(%T).GetPackage() = %q; want %q", enum, pkg, want)
+	}
+	if name, want := ed.GetName(), "Type"; name != want {
+		t.Errorf("descriptor.ForEnum(%T).GetName() = %q; want %q", enum, name, want)
+	}
+	if value, want := ed.GetValue()[0].GetName(), "TYPE_DOUBLE"; value != want {
+		t.Errorf("descriptor.ForEnum(%T).GetValue()[0].GetName() = %v; want %v", enum, value, want)
+	}
+}
+
 func Example_Options() {
 	var msg *tpb.MyMessageSet
 	_, md := descriptor.ForMessage(msg)
@@ -30,3 +44,14 @@ func Example_Options() {
 	// Output:
 	// MyMessageSet uses option message_set_wire_format.
 }
+
+func Example_EnumOptions() {
+	var enum *tpb.EnumAllowingAlias_NUMBER
+	_, ed := descriptor.ForEnum(enum)
+	if ed.GetOptions().GetAllowAlias() {
+		fmt.Printf("%v uses option allow_alias.\n", ed.GetName())
+	}
+
+	// Output:
+	// NUMBER uses option allow_alias.
+}