Skip to content

Handle pickling top-level enums. #34568

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tvalentyn merged 4 commits into from
Apr 8, 2025
Merged

Handle pickling top-level enums. #34568

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
badf0fc
Handle pickling top-level enums.
claudevdm Apr 7, 2025
ad78dd6
Add tests.
claudevdm Apr 7, 2025
28542a9
Lint fix.
Apr 8, 2025
9b24b0f
Add license header.
Apr 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions sdks/java/extensions/protobuf/src/test/proto/proto2_coder_test_messages.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
68 changes: 41 additions & 27 deletions sdks/python/apache_beam/coders/proto2_coder_test_messages_pb2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
5 changes: 5 additions & 0 deletions sdks/python/apache_beam/internal/cloudpickle_pickler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
16 changes: 12 additions & 4 deletions sdks/python/apache_beam/internal/cloudpickle_pickler_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))())

Expand Down
Loading