diff --git a/python/google/protobuf/internal/enum_type_wrapper.py b/python/google/protobuf/internal/enum_type_wrapper.py index cc65fc0bc865..1fa6ab9f5803 100644 --- a/python/google/protobuf/internal/enum_type_wrapper.py +++ b/python/google/protobuf/internal/enum_type_wrapper.py @@ -12,6 +12,8 @@ reflection_test.py """ +import sys + __author__ = 'rabsatt@google.com (Kevin Rabsatt)' @@ -99,3 +101,12 @@ def __getattr__(self, name): pass # fall out to break exception chaining raise AttributeError('Enum {} has no value defined for name {!r}'.format( self._enum_type.name, name)) + + def __or__(self, other): + """Returns the union type of self and other.""" + if sys.version_info >= (3, 10): + return type(self) | other + else: + raise NotImplementedError( + 'You may not use | on EnumTypes (or classes) below python 3.10' + ) diff --git a/python/google/protobuf/internal/message_test.py b/python/google/protobuf/internal/message_test.py index 68f1999b1671..7b4478e24586 100755 --- a/python/google/protobuf/internal/message_test.py +++ b/python/google/protobuf/internal/message_test.py @@ -22,6 +22,7 @@ import pickle import pydoc import sys +import types import unittest from unittest import mock import warnings @@ -30,6 +31,7 @@ from google.protobuf.internal import api_implementation # pylint: disable=g-import-not-at-top from google.protobuf.internal import encoder +from google.protobuf.internal import enum_type_wrapper from google.protobuf.internal import more_extensions_pb2 from google.protobuf.internal import more_messages_pb2 from google.protobuf.internal import packed_field_test_pb2 @@ -1314,6 +1316,26 @@ def __eq__(self, other): self.assertNotEqual(m, ComparesWithFoo()) self.assertNotEqual(ComparesWithFoo(), m) + def testTypeUnion(self, message_module): + # Below python 3.10 you cannot create union types with the | operator, so we + # skip testing for unions with old versions. + if sys.version_info < (3, 10): + return + enum_type = enum_type_wrapper.EnumTypeWrapper( + message_module.TestAllTypes.NestedEnum.DESCRIPTOR + ) + union_type = enum_type | int + self.assertIsInstance(union_type, types.UnionType) + + def get_union() -> union_type: + return enum_type + + union = get_union() + self.assertIsInstance(union, enum_type_wrapper.EnumTypeWrapper) + self.assertEqual( + union.DESCRIPTOR, message_module.TestAllTypes.NestedEnum.DESCRIPTOR + ) + # Class to test proto2-only features (required, extensions, etc.) @testing_refleaks.TestCase