Skip to content

Commit 0324f39

Browse files
anandoleecopybara-github
authored andcommitted
Add clear() method to repeated fields in Python.
Example usage: msg.repeated_int32.extend([1,2]) msg.repeated_int32.clear() PiperOrigin-RevId: 736221526
1 parent fdb9018 commit 0324f39

File tree

6 files changed

+120
-29
lines changed

6 files changed

+120
-29
lines changed

python/google/protobuf/internal/message_test.py

+34
Original file line numberDiff line numberDiff line change
@@ -498,6 +498,40 @@ def testMergeFromMissingRequiredField(self, message_module):
498498
message.MergeFrom(msg)
499499
self.assertEqual(msg, message)
500500

501+
def testScalarRepeatedClear(self, message_module):
502+
msg = message_module.TestAllTypes()
503+
empty_size = msg.ByteSize()
504+
msg.repeated_int32.append(1)
505+
msg.repeated_int32.append(3)
506+
repeated_int = msg.repeated_int32
507+
self.assertEqual(2, len(msg.repeated_int32))
508+
self.assertGreater(msg.ByteSize(), empty_size)
509+
msg.repeated_int32.clear()
510+
self.assertEqual(0, len(msg.repeated_int32))
511+
self.assertEqual(0, len(repeated_int))
512+
self.assertEqual(empty_size, msg.ByteSize())
513+
514+
def testCompositeRepeatedClear(self, message_module):
515+
msg = message_module.TestAllTypes()
516+
empty_size = msg.ByteSize()
517+
msg.repeated_nested_message.add(bb=123)
518+
msg.repeated_nested_message.add(bb=2)
519+
repeated_nested_message = msg.repeated_nested_message
520+
self.assertEqual(2, len(msg.repeated_nested_message))
521+
self.assertGreater(msg.ByteSize(), empty_size)
522+
msg.repeated_nested_message.clear()
523+
self.assertEqual(0, len(msg.repeated_nested_message))
524+
self.assertEqual(0, len(repeated_nested_message))
525+
self.assertEqual(empty_size, msg.ByteSize())
526+
527+
def testCompositeRepeatedClearRelease(self, message_module):
528+
msg = message_module.TestAllTypes()
529+
msg.repeated_nested_message.add(bb=123)
530+
# sub msg reference should still work after clear()
531+
sub_msg = msg.repeated_nested_message[0]
532+
msg.repeated_nested_message.clear()
533+
self.assertEqual(123, sub_msg.bb)
534+
501535
def testAddWrongRepeatedNestedField(self, message_module):
502536
msg = message_module.TestAllTypes()
503537
try:

python/google/protobuf/pyext/message.cc

+42-29
Original file line numberDiff line numberDiff line change
@@ -851,6 +851,42 @@ static PyObject* GetIntegerEnumValue(const FieldDescriptor& descriptor,
851851
return value;
852852
}
853853

854+
void DeleteLastRepeatedWithSize(CMessage* self,
855+
const FieldDescriptor* field_descriptor,
856+
Py_ssize_t n) {
857+
Message* message = self->message;
858+
const Reflection* reflection = message->GetReflection();
859+
ABSL_DCHECK(reflection->FieldSize(*message, field_descriptor) >= n);
860+
Arena* arena = message->GetArena();
861+
ABSL_DCHECK_EQ(arena, nullptr)
862+
<< "python protobuf is expected to be allocated from heap";
863+
for (; n > 0; n--) {
864+
// It seems that RemoveLast() is less efficient for sub-messages, and
865+
// the memory is not completely released. Prefer ReleaseLast().
866+
//
867+
// To work around a debug hardening (PROTOBUF_FORCE_COPY_IN_RELEASE),
868+
// explicitly use UnsafeArenaReleaseLast. To not break rare use cases where
869+
// arena is used, we fallback to ReleaseLast (but ABSL_DCHECK to find/fix
870+
// it).
871+
//
872+
// Note that arena is likely null and ABSL_DCHECK and ReleaseLast might be
873+
// redundant. The current approach takes extra cautious path not to disrupt
874+
// production.
875+
Message* sub_message =
876+
(arena == nullptr)
877+
? reflection->UnsafeArenaReleaseLast(message, field_descriptor)
878+
: reflection->ReleaseLast(message, field_descriptor);
879+
// If there is a live weak reference to an item being removed, we "Release"
880+
// it, and it takes ownership of the message.
881+
if (CMessage* released = self->MaybeReleaseSubMessage(sub_message)) {
882+
released->message = sub_message;
883+
} else {
884+
// sub_message was not transferred, delete it.
885+
delete sub_message;
886+
}
887+
}
888+
}
889+
854890
// Delete a slice from a repeated field.
855891
// The only way to remove items in C++ protos is to delete the last one,
856892
// so we swap items to move the deleted ones at the end, and then strip the
@@ -911,37 +947,14 @@ int DeleteRepeatedField(CMessage* self, const FieldDescriptor* field_descriptor,
911947
}
912948
}
913949

914-
Arena* arena = message->GetArena();
915-
ABSL_DCHECK_EQ(arena, nullptr)
916-
<< "python protobuf is expected to be allocated from heap";
917950
// Remove items, starting from the end.
918-
for (; length > to; length--) {
919-
if (field_descriptor->cpp_type() != FieldDescriptor::CPPTYPE_MESSAGE) {
951+
if (field_descriptor->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
952+
DeleteLastRepeatedWithSize(self, field_descriptor, length - to);
953+
} else {
954+
ABSL_DCHECK_EQ(message->GetArena(), nullptr)
955+
<< "python protobuf is expected to be allocated from heap";
956+
for (; length > to; length--) {
920957
reflection->RemoveLast(message, field_descriptor);
921-
continue;
922-
}
923-
// It seems that RemoveLast() is less efficient for sub-messages, and
924-
// the memory is not completely released. Prefer ReleaseLast().
925-
//
926-
// To work around a debug hardening (PROTOBUF_FORCE_COPY_IN_RELEASE),
927-
// explicitly use UnsafeArenaReleaseLast. To not break rare use cases where
928-
// arena is used, we fallback to ReleaseLast (but ABSL_DCHECK to find/fix
929-
// it).
930-
//
931-
// Note that arena is likely null and ABSL_DCHECK and ReleaseLast might be
932-
// redundant. The current approach takes extra cautious path not to disrupt
933-
// production.
934-
Message* sub_message =
935-
(arena == nullptr)
936-
? reflection->UnsafeArenaReleaseLast(message, field_descriptor)
937-
: reflection->ReleaseLast(message, field_descriptor);
938-
// If there is a live weak reference to an item being removed, we "Release"
939-
// it, and it takes ownership of the message.
940-
if (CMessage* released = self->MaybeReleaseSubMessage(sub_message)) {
941-
released->message = sub_message;
942-
} else {
943-
// sub_message was not transferred, delete it.
944-
delete sub_message;
945958
}
946959
}
947960

python/google/protobuf/pyext/message.h

+5
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,11 @@ const FieldDescriptor* GetExtensionDescriptor(PyObject* extension);
165165
CMessage* InternalGetSubMessage(CMessage* self,
166166
const FieldDescriptor* field_descriptor);
167167

168+
// Delete the last n items in a repeated field.
169+
void DeleteLastRepeatedWithSize(CMessage* self,
170+
const FieldDescriptor* field_descriptor,
171+
Py_ssize_t n);
172+
168173
// Deletes a range of items in a repeated field (following a
169174
// removal in a RepeatedCompositeContainer).
170175
//

python/google/protobuf/pyext/repeated_composite_container.cc

+13
Original file line numberDiff line numberDiff line change
@@ -418,6 +418,17 @@ static PyObject* Reverse(PyObject* pself) {
418418
}
419419

420420
// ---------------------------------------------------------------------
421+
static PyObject* Clear(PyObject* pself) {
422+
RepeatedCompositeContainer* self =
423+
reinterpret_cast<RepeatedCompositeContainer*>(pself);
424+
CMessage* cmessage = self->parent;
425+
Message* message = cmessage->message;
426+
const FieldDescriptor* field_descriptor = self->parent_field_descriptor;
427+
const Reflection* reflection = message->GetReflection();
428+
Py_ssize_t length = reflection->FieldSize(*message, field_descriptor);
429+
cmessage::DeleteLastRepeatedWithSize(cmessage, field_descriptor, length);
430+
Py_RETURN_NONE;
431+
}
421432

422433
static PyObject* Item(PyObject* pself, Py_ssize_t index) {
423434
RepeatedCompositeContainer* self =
@@ -511,6 +522,8 @@ static PyMethodDef Methods[] = {
511522
"Sorts the repeated container."},
512523
{"reverse", reinterpret_cast<PyCFunction>(Reverse), METH_NOARGS,
513524
"Reverses elements order of the repeated container."},
525+
{"clear", reinterpret_cast<PyCFunction>(Clear), METH_NOARGS,
526+
"Clears the repeated container."},
514527
{"MergeFrom", MergeFromMethod, METH_O,
515528
"Adds objects to the repeated container."},
516529
{nullptr, nullptr}};

python/google/protobuf/pyext/repeated_scalar_container.cc

+13
Original file line numberDiff line numberDiff line change
@@ -599,6 +599,17 @@ static PyObject* Reverse(PyObject* pself) {
599599
Py_RETURN_NONE;
600600
}
601601

602+
static PyObject* Clear(PyObject* pself) {
603+
RepeatedScalarContainer* self =
604+
reinterpret_cast<RepeatedScalarContainer*>(pself);
605+
CMessage* cmessage = self->parent;
606+
cmessage::AssureWritable(cmessage);
607+
Message* message = cmessage->message;
608+
const FieldDescriptor* field_descriptor = self->parent_field_descriptor;
609+
message->GetReflection()->ClearField(message, field_descriptor);
610+
Py_RETURN_NONE;
611+
}
612+
602613
static PyObject* Pop(PyObject* pself, PyObject* args) {
603614
Py_ssize_t index = -1;
604615
if (!PyArg_ParseTuple(args, "|n", &index)) {
@@ -693,6 +704,8 @@ static PyMethodDef Methods[] = {
693704
"Sorts the repeated container."},
694705
{"reverse", reinterpret_cast<PyCFunction>(Reverse), METH_NOARGS,
695706
"Reverses elements order of the repeated container."},
707+
{"clear", reinterpret_cast<PyCFunction>(Clear), METH_NOARGS,
708+
"Clears the repeated container."},
696709
{"MergeFrom", static_cast<PyCFunction>(MergeFrom), METH_O,
697710
"Merges a repeated container into the current container."},
698711
{nullptr, nullptr}};

python/repeated.c

+13
Original file line numberDiff line numberDiff line change
@@ -530,6 +530,15 @@ static PyObject* PyUpb_RepeatedContainer_Reverse(PyObject* _self) {
530530
Py_RETURN_NONE;
531531
}
532532

533+
static PyObject* PyUpb_RepeatedContainer_Clear(PyObject* _self) {
534+
PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
535+
Py_ssize_t size = PyUpb_RepeatedContainer_Length(_self);
536+
if (size > 0) {
537+
upb_Array_Delete(self->ptr.arr, 0, size);
538+
}
539+
Py_RETURN_NONE;
540+
}
541+
533542
static PyObject* PyUpb_RepeatedContainer_MergeFrom(PyObject* _self,
534543
PyObject* args) {
535544
return PyUpb_RepeatedContainer_Extend(_self, args);
@@ -638,6 +647,8 @@ static PyMethodDef PyUpb_RepeatedCompositeContainer_Methods[] = {
638647
METH_VARARGS | METH_KEYWORDS, "Sorts the repeated container."},
639648
{"reverse", (PyCFunction)PyUpb_RepeatedContainer_Reverse, METH_NOARGS,
640649
"Reverses elements order of the repeated container."},
650+
{"clear", (PyCFunction)PyUpb_RepeatedContainer_Clear, METH_NOARGS,
651+
"Clears repeated container."},
641652
{"MergeFrom", PyUpb_RepeatedContainer_MergeFrom, METH_O,
642653
"Adds objects to the repeated container."},
643654
{NULL, NULL}};
@@ -734,6 +745,8 @@ static PyMethodDef PyUpb_RepeatedScalarContainer_Methods[] = {
734745
METH_VARARGS | METH_KEYWORDS, "Sorts the repeated container."},
735746
{"reverse", (PyCFunction)PyUpb_RepeatedContainer_Reverse, METH_NOARGS,
736747
"Reverses elements order of the repeated container."},
748+
{"clear", (PyCFunction)PyUpb_RepeatedContainer_Clear, METH_NOARGS,
749+
"Clears repeated container."},
737750
{"MergeFrom", PyUpb_RepeatedContainer_MergeFrom, METH_O,
738751
"Merges a repeated container into the current container."},
739752
{NULL, NULL}};

0 commit comments

Comments
 (0)