feat: allow_alias for enums

.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