feat: allow_alias for enums (#207)

Certain APIs, e.g. recommendationengine, have enums where variants are
aliased, i.e. different names map to the same integer value.

Allowing this behavior in proto plus for the cpp protobuf runtime
requires constructing and passing the correct options.

.github/sync-repo-settings.yaml

@@ -12,7 +12,7 @@ branchProtectionRules:
     - 'unit (3.7)'
     - 'unit (3.7, cpp)'
     - 'unit (3.8)'
-    - 'unit (3.9, cpp)'
+    # - 'unit (3.9, cpp)' # Don't have binary wheels for 3.9 cpp protobuf yet
     - 'unit (3.9)'
     - 'cla/google'
   requiredApprovingReviewCount: 1

proto/enums.py

@@ -47,6 +47,21 @@ def __new__(mcls, name, bases, attrs):
         filename = _file_info._FileInfo.proto_file_name(
             attrs.get("__module__", name.lower())
         )
+
+        # Retrieve any enum options.
+        # We expect something that looks like an EnumOptions message,
+        # either an actual instance or a dict-like representation.
+        pb_options = "_pb_options"
+        opts = attrs.pop(pb_options, {})
+        # This is the only portable way to remove the _pb_options name
+        # from the enum attrs.
+        # In 3.7 onwards, we can define an _ignore_ attribute and do some
+        # mucking around with that.
+        if pb_options in attrs._member_names:
+            idx = attrs._member_names.index(pb_options)
+            attrs._member_names.pop(idx)
+
+        # Make the descriptor.
         enum_desc = descriptor_pb2.EnumDescriptorProto(
             name=name,
             # Note: the superclass ctor removes the variants, so get them now.
@@ -60,6 +75,7 @@ def __new__(mcls, name, bases, attrs):
                 ),
                 key=lambda v: v.number,
             ),
+            options=opts,
         )
 
         file_info = _file_info._FileInfo.maybe_add_descriptor(filename, package)

tests/test_fields_enum.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import proto
+import pytest
 import sys
 
 
@@ -353,3 +355,40 @@ class Task(proto.Message):
     t = Task(weekday="TUESDAY")
     t2 = Task.deserialize(Task.serialize(t))
     assert t == t2
+
+
+if os.environ.get("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION", "python") == "cpp":
+    # This test only works, and is only relevant, with the cpp runtime.
+    # Python just doesn't give a care and lets it work anyway.
+    def test_enum_alias_bad():
+        # Certain enums may shadow the different enum monikers with the same value.
+        # This is generally discouraged, and protobuf will object by default,
+        # but will explicitly allow this behavior if the enum is defined with
+        # the `allow_alias` option set.
+        with pytest.raises(TypeError):
+
+            # The wrapper message is a hack to avoid manifest wrangling to
+            # define the enum.
+            class BadMessage(proto.Message):
+                class BadEnum(proto.Enum):
+                    UNKNOWN = 0
+                    DEFAULT = 0
+
+                bad_dup_enum = proto.Field(proto.ENUM, number=1, enum=BadEnum)
+
+
+def test_enum_alias_good():
+    # Have to split good and bad enum alias into two tests so that the generated
+    # file descriptor is properly created.
+    # For the python runtime, aliases are allowed by default, but we want to
+    # make sure that the options don't cause problems.
+    # For the cpp runtime, we need to verify that we can in fact define aliases.
+    class GoodMessage(proto.Message):
+        class GoodEnum(proto.Enum):
+            _pb_options = {"allow_alias": True}
+            UNKNOWN = 0
+            DEFAULT = 0
+
+        good_dup_enum = proto.Field(proto.ENUM, number=1, enum=GoodEnum)
+
+    assert GoodMessage.GoodEnum.UNKNOWN == GoodMessage.GoodEnum.DEFAULT == 0