Nextgen Proto Pythonic API: Add byte_size(), clear_message() and clear_field() APIs in proto module

PiperOrigin-RevId: 714287607
This commit is contained in:
Jie Luo 2025-01-10 17:57:34 -08:00 committed by Copybara-Service
parent 5b0fa1d84e
commit 975bb07baa
2 changed files with 59 additions and 1 deletions

View file

@ -77,6 +77,27 @@ class ProtoTest(unittest.TestCase):
str(context.exception),
)
def test_byte_size(self, message_module):
msg = message_module.TestAllTypes()
self.assertEqual(0, proto.byte_size(msg))
msg.optional_int32 = 123
self.assertEqual(2, proto.byte_size(msg))
def test_clear_message(self, message_module):
msg = message_module.TestAllTypes()
msg.oneof_uint32 = 11
msg.repeated_nested_message.add(bb=1)
proto.clear_message(msg)
self.assertIsNone(msg.WhichOneof('oneof_field'))
self.assertEqual(0, len(msg.repeated_nested_message))
def test_clear_field(self, message_module):
msg = message_module.TestAllTypes()
msg.optional_int32 = 123
self.assertEqual(123, msg.optional_int32)
proto.clear_field(msg, 'optional_int32')
self.assertEqual(0, msg.optional_int32)
class SelfFieldTest(unittest.TestCase):

View file

@ -8,7 +8,7 @@
"""Contains the Nextgen Pythonic protobuf APIs."""
import io
from typing import Type, TypeVar
from typing import Text, Type, TypeVar
from google.protobuf.internal import decoder
from google.protobuf.internal import encoder
@ -114,3 +114,40 @@ def parse_length_prefixed(
'{2}.'.format(size, parsed_size, message.DESCRIPTOR.name)
)
return message
def byte_size(message: Message) -> int:
"""Returns the serialized size of this message.
Args:
message: A proto message.
Returns:
int: The number of bytes required to serialize this message.
"""
return message.ByteSize()
def clear_message(message: Message) -> None:
"""Clears all data that was set in the message.
Args:
message: The proto message to be cleared.
"""
message.Clear()
def clear_field(message: Message, field_name: Text) -> None:
"""Clears the contents of a given field.
Inside a oneof group, clears the field set. If the name neither refers to a
defined field or oneof group, :exc:`ValueError` is raised.
Args:
message: The proto message.
field_name (str): The name of the field to be cleared.
Raises:
ValueError: if the `field_name` is not a member of this message.
"""
message.ClearField(field_name)