How to match protobuf enum constants with values in Python

It is sometimes very useful to be able to enumerate, in Python, all the values of an enum defined in a protobuf schema. For example, your code might receive a number representing an enum value, and you want to print a human-readable representation of it. While Google's official documentation states that "No type corresponding to EnumName is defined" in the generated code 1, we can still achieve this functionality by way of undocumented implementation details. Normally, it's a good idea to not rely on implementation details in your code, since they could be changed from underneath you, so this technique is probably best left to quick testing scripts. (I did in fact use this technique in a testing script, which inspired this article.)

Let's say you have an enum defined as follows:

enum MyEnumType {
    ENUM_VALUE_A = 100;
    ENUM_VALUE_B = 120;
    ENUM_VALUE_C = 140;
}

And you want a function that does a "reverse lookup" from the enum value to the enum name, like so:

def enum_value_to_name(val):
    # ...

enum_value_to_name(140) -> "ENUM_VALUE_C"
enum_value_to_name(100) -> "ENUM_VALUE_A"
enum_value_to_name(120) -> "ENUM_VALUE_B"

Here I'll present a simple implementation of enum_value_to_name, and afterwards how I arrived at this implementation:

def enum_value_to_name(val):
    desc = MyEnumType.DESCRIPTOR
    for (k,v) in desc.values_by_name.items():
        if v.number == val:
            return k
    return None # if val isn't a value in MyEnumType

This naive implementation is pretty easy to come by if you use do a little Python sleuthing. Compile the .proto file and import that module in a Python shell, for example:

$ python
>>> import myproto_pb2 as pb2
>>> pb2.MyEnumType
<google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper object at 0x10cfb6d10>

        (as we can see, a type *does* get defined for MyEnumType,
         but it's an internal implementation detail)

>>> pb2.MyEnumType.__dict__
{'DESCRIPTOR': <google.protobuf.descriptor.EnumDescriptor object at 0x10cfb6cd0>, '_enum_type': <google.protobuf.descriptor.EnumDescriptor object at 0x10cfb6cd0>}

        (__dict__ is a dictionary representation of all the
         attributes that an object has. It's an easy way to
         inspect the internals of an object. 'DESCRIPTOR'
         looks interesting, let's see what is in there)

>>> pb2.MyEnumType.DESCRIPTOR.__dict__
{'_serialized_end': 48, '_options_class_name': 'EnumOptions', 'name': 'MyEnumType', '_options': None, 'full_name': 'MyEnumType', 'containing_type': None, 'values': [<google.protobuf.descriptor.EnumValueDescriptor object at 0x10cfb6c50>], 'values_by_name': {'ENUM_VALUE_A': <google.protobuf.descriptor.EnumValueDescriptor object at 0x10cfb6c50>}, 'file': <google.protobuf.descriptor.FileDescriptor object at 0x10cfb6c10>, 'has_options': False, '_serialized_start': 18, 'values_by_number': {100: <google.protobuf.descriptor.EnumValueDescriptor object at 0x10cfb6c50>}}

        (ooh, 'values_by_name' looks like what we want. Now
         all we need to do is to extract a number (integer)
         from those EnumValueDescriptor objects)

>>> desc = pb2.MyEnumType.DESCRIPTOR
>>> desc.values_by_name['ENUM_VALUE_A'].__dict__
{'index': 0, '_options_class_name': 'EnumValueOptions', 'name': 'ENUM_VALUE_A', '_options': None, 'type': <google.protobuf.descriptor.EnumDescriptor object at 0x10cfb6cd0>, 'number': 100, 'has_options': False}

        (nice, looks like the 'number' attribute is what
         we want)

>>> for k,v in desc.values_by_name.items():
...   print("{} -> {}".format(k, v.number))
ENUM_VALUE_A -> 100
ENUM_VALUE_B -> 120
ENUM_VALUE_C -> 140

And there you have it, after a bit of digging into the internals of the protobuf generated code, we can easily get a dictionary containing the names (as strings) and values (as numbers) of all the entries in the enum.

Comments

comments powered by Disqus