Handle pickling top-level enums. (#34568)

* Handle pickling top-level enums.

* Add tests.

* Lint fix.

* Add license header.

---------

Co-authored-by: Claude <[email protected]>

Author: claudevdm

Diff for: sdks/java/extensions/protobuf/src/test/proto/proto2_coder_test_messages.proto

@@ -51,3 +51,17 @@ message MessageWithMap {
 message ReferencesMessageWithMap {
   repeated MessageWithMap field1 = 1;
 }
+
+enum TopLevelEnum {
+  UNSPECIFIED = 0;
+  ONE = 1;
+  TWO = 2;
+}
+
+message MessageD {
+  enum NestedEnum {
+    UNSPECIFIED = 0;
+    ONE = 1;
+    TWO = 2;
+  }
+}

Diff for: sdks/python/apache_beam/coders/proto2_coder_test_messages_pb2.py

@@ -17,42 +17,56 @@
 
 # -*- coding: utf-8 -*-
 # Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
 # source: apache_beam/coders/proto2_coder_test_messages.proto
-
+# Protobuf Python Version: 5.28.0
 """Generated protocol buffer code."""
-from google.protobuf.internal import builder as _builder
 from google.protobuf import descriptor as _descriptor
 from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
 from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    5,
+    28,
+    0,
+    '',
+    'apache_beam/coders/proto2_coder_test_messages.proto'
+)
 # @@protoc_insertion_point(imports)
 
 _sym_db = _symbol_database.Default()
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
-    b'\n3apache_beam/coders/proto2_coder_test_messages.proto\x12\x1aproto2_coder_test_messages\"P\n\x08MessageA\x12\x0e\n\x06\x66ield1\x18\x01 \x01(\t\x12\x34\n\x06\x66ield2\x18\x02 \x03(\x0b\x32$.proto2_coder_test_messages.MessageB\"\x1a\n\x08MessageB\x12\x0e\n\x06\x66ield1\x18\x01 \x01(\x08\"\x10\n\x08MessageC*\x04\x08\x64\x10j\"\xad\x01\n\x0eMessageWithMap\x12\x46\n\x06\x66ield1\x18\x01 \x03(\x0b\x32\x36.proto2_coder_test_messages.MessageWithMap.Field1Entry\x1aS\n\x0b\x46ield1Entry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x33\n\x05value\x18\x02 \x01(\x0b\x32$.proto2_coder_test_messages.MessageA:\x02\x38\x01\"V\n\x18ReferencesMessageWithMap\x12:\n\x06\x66ield1\x18\x01 \x03(\x0b\x32*.proto2_coder_test_messages.MessageWithMap:Z\n\x06\x66ield1\x12$.proto2_coder_test_messages.MessageC\x18\x65 \x01(\x0b\x32$.proto2_coder_test_messages.MessageA:Z\n\x06\x66ield2\x12$.proto2_coder_test_messages.MessageC\x18\x66 \x01(\x0b\x32$.proto2_coder_test_messages.MessageBB)\n\'org.apache.beam.sdk.extensions.protobuf'
-)
 
-_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
-_builder.BuildTopDescriptorsAndMessages(
-    DESCRIPTOR, 'apache_beam.coders.proto2_coder_test_messages_pb2', globals())
-if _descriptor._USE_C_DESCRIPTORS == False:
-  MessageC.RegisterExtension(field1)
-  MessageC.RegisterExtension(field2)
 
-  DESCRIPTOR._options = None
-  DESCRIPTOR._serialized_options = b'\n\'org.apache.beam.sdk.extensions.protobuf'
-  _MESSAGEWITHMAP_FIELD1ENTRY._options = None
-  _MESSAGEWITHMAP_FIELD1ENTRY._serialized_options = b'8\001'
-  _MESSAGEA._serialized_start = 83
-  _MESSAGEA._serialized_end = 163
-  _MESSAGEB._serialized_start = 165
-  _MESSAGEB._serialized_end = 191
-  _MESSAGEC._serialized_start = 193
-  _MESSAGEC._serialized_end = 209
-  _MESSAGEWITHMAP._serialized_start = 212
-  _MESSAGEWITHMAP._serialized_end = 385
-  _MESSAGEWITHMAP_FIELD1ENTRY._serialized_start = 302
-  _MESSAGEWITHMAP_FIELD1ENTRY._serialized_end = 385
-  _REFERENCESMESSAGEWITHMAP._serialized_start = 387
-  _REFERENCESMESSAGEWITHMAP._serialized_end = 473
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n3apache_beam/coders/proto2_coder_test_messages.proto\x12\x1aproto2_coder_test_messages\"P\n\x08MessageA\x12\x0e\n\x06\x66ield1\x18\x01 \x01(\t\x12\x34\n\x06\x66ield2\x18\x02 \x03(\x0b\x32$.proto2_coder_test_messages.MessageB\"\x1a\n\x08MessageB\x12\x0e\n\x06\x66ield1\x18\x01 \x01(\x08\"\x10\n\x08MessageC*\x04\x08\x64\x10j\"\xad\x01\n\x0eMessageWithMap\x12\x46\n\x06\x66ield1\x18\x01 \x03(\x0b\x32\x36.proto2_coder_test_messages.MessageWithMap.Field1Entry\x1aS\n\x0b\x46ield1Entry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x33\n\x05value\x18\x02 \x01(\x0b\x32$.proto2_coder_test_messages.MessageA:\x02\x38\x01\"V\n\x18ReferencesMessageWithMap\x12:\n\x06\x66ield1\x18\x01 \x03(\x0b\x32*.proto2_coder_test_messages.MessageWithMap\";\n\x08MessageD\"/\n\nNestedEnum\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x07\n\x03ONE\x10\x01\x12\x07\n\x03TWO\x10\x02*1\n\x0cTopLevelEnum\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x07\n\x03ONE\x10\x01\x12\x07\n\x03TWO\x10\x02:Z\n\x06\x66ield1\x12$.proto2_coder_test_messages.MessageC\x18\x65 \x01(\x0b\x32$.proto2_coder_test_messages.MessageA:Z\n\x06\x66ield2\x12$.proto2_coder_test_messages.MessageC\x18\x66 \x01(\x0b\x32$.proto2_coder_test_messages.MessageBB)\n\'org.apache.beam.sdk.extensions.protobuf')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'apache_beam.coders.proto2_coder_test_messages_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'\n\'org.apache.beam.sdk.extensions.protobuf'
+  _globals['_MESSAGEWITHMAP_FIELD1ENTRY']._loaded_options = None
+  _globals['_MESSAGEWITHMAP_FIELD1ENTRY']._serialized_options = b'8\001'
+  _globals['_TOPLEVELENUM']._serialized_start=536
+  _globals['_TOPLEVELENUM']._serialized_end=585
+  _globals['_MESSAGEA']._serialized_start=83
+  _globals['_MESSAGEA']._serialized_end=163
+  _globals['_MESSAGEB']._serialized_start=165
+  _globals['_MESSAGEB']._serialized_end=191
+  _globals['_MESSAGEC']._serialized_start=193
+  _globals['_MESSAGEC']._serialized_end=209
+  _globals['_MESSAGEWITHMAP']._serialized_start=212
+  _globals['_MESSAGEWITHMAP']._serialized_end=385
+  _globals['_MESSAGEWITHMAP_FIELD1ENTRY']._serialized_start=302
+  _globals['_MESSAGEWITHMAP_FIELD1ENTRY']._serialized_end=385
+  _globals['_REFERENCESMESSAGEWITHMAP']._serialized_start=387
+  _globals['_REFERENCESMESSAGEWITHMAP']._serialized_end=473
+  _globals['_MESSAGED']._serialized_start=475
+  _globals['_MESSAGED']._serialized_end=534
+  _globals['_MESSAGED_NESTEDENUM']._serialized_start=487
+  _globals['_MESSAGED_NESTEDENUM']._serialized_end=534
 # @@protoc_insertion_point(module_scope)

Diff for: sdks/python/apache_beam/internal/cloudpickle_pickler.py

@@ -86,6 +86,11 @@ def _reconstruct_enum_descriptor(full_name):
     if not hasattr(module, 'DESCRIPTOR'):
       continue
 
+    if hasattr(module.DESCRIPTOR, 'enum_types_by_name'):
+      for (_, enum_desc) in module.DESCRIPTOR.enum_types_by_name.items():
+        if enum_desc.full_name == full_name:
+          return enum_desc
+
     for _, attr_value in vars(module).items():
       if not hasattr(attr_value, 'DESCRIPTOR'):
         continue

Diff for: sdks/python/apache_beam/internal/cloudpickle_pickler_test.py

@@ -23,21 +23,29 @@
 import types
 import unittest
 
+from apache_beam.coders import proto2_coder_test_messages_pb2
 from apache_beam.internal import module_test
 from apache_beam.internal.cloudpickle_pickler import dumps
 from apache_beam.internal.cloudpickle_pickler import loads
-from apache_beam.portability.api import beam_runner_api_pb2
 
 
 class PicklerTest(unittest.TestCase):
 
   NO_MAPPINGPROXYTYPE = not hasattr(types, "MappingProxyType")
 
-  def test_pickle_enum_descriptor(self):
-    TimeDomain = beam_runner_api_pb2.TimeDomain.Enum
+  def test_pickle_nested_enum_descriptor(self):
+    NestedEnum = proto2_coder_test_messages_pb2.MessageD.NestedEnum
 
     def fn():
-      return TimeDomain.EVENT_TIME
+      return NestedEnum.TWO
+
+    self.assertEqual(fn(), loads(dumps(fn))())
+
+  def test_pickle_top_level_enum_descriptor(self):
+    TopLevelEnum = proto2_coder_test_messages_pb2.TopLevelEnum
+
+    def fn():
+      return TopLevelEnum.ONE
 
     self.assertEqual(fn(), loads(dumps(fn))())