From 9c76ea4413269fa344d6b586d478925b6ed014e7 Mon Sep 17 00:00:00 2001 From: Jie Luo Date: Wed, 30 Aug 2023 15:45:27 -0700 Subject: [PATCH] Add upb python Py_mp_subscript for descriptor sequences. A behavior diff that between upb and the other two implementations. PiperOrigin-RevId: 561471616 --- .../protobuf/internal/descriptor_test.py | 9 +++++ upb/python/descriptor_containers.c | 23 +++++++++++- upb/python/protobuf.c | 28 ++++++++++++++ upb/python/protobuf.h | 7 ++++ upb/python/repeated.c | 37 +------------------ 5 files changed, 68 insertions(+), 36 deletions(-) diff --git a/python/google/protobuf/internal/descriptor_test.py b/python/google/protobuf/internal/descriptor_test.py index 43ce965e7d..dda15bde66 100755 --- a/python/google/protobuf/internal/descriptor_test.py +++ b/python/google/protobuf/internal/descriptor_test.py @@ -810,6 +810,15 @@ class GeneratedDescriptorTest(unittest.TestCase): oneof_descriptor.full_name) self.assertEqual(0, oneof_descriptor.index) + def testDescriptorSlice(self): + message_descriptor = unittest_pb2.TestAllTypes.DESCRIPTOR + nested = message_descriptor.nested_types[:] + self.assertEqual(message_descriptor.nested_types, nested) + fields = message_descriptor.fields + fields_list = list(fields) + self.assertEqual(fields_list[:], fields[:]) + self.assertEqual(fields_list[2::2], fields[2::2]) + self.assertEqual(fields_list[3:19:3], fields[3:19:3]) class DescriptorCopyToProtoTest(unittest.TestCase): """Tests for CopyTo functions of Descriptor.""" diff --git a/upb/python/descriptor_containers.c b/upb/python/descriptor_containers.c index 3c959b4a58..e1eacb2a77 100644 --- a/upb/python/descriptor_containers.c +++ b/upb/python/descriptor_containers.c @@ -210,6 +210,9 @@ static PyObject* PyUpb_GenericSequence_GetItem(PyObject* _self, Py_ssize_t index) { PyUpb_GenericSequence* self = PyUpb_GenericSequence_Self(_self); Py_ssize_t size = self->funcs->get_elem_count(self->parent); + if (index < 0) { + index += size; + } if (index < 0 || index >= size) { PyErr_Format(PyExc_IndexError, "list index (%zd) out of range", index); return NULL; @@ -267,6 +270,24 @@ static PyObject* PyUpb_GenericSequence_RichCompare(PyObject* _self, return PyBool_FromLong(ret); } +static PyObject* PyUpb_GenericSequence_Subscript(PyObject* _self, + PyObject* item) { + PyUpb_GenericSequence* self = PyUpb_GenericSequence_Self(_self); + Py_ssize_t size = self->funcs->get_elem_count(self->parent); + Py_ssize_t idx, count, step; + if (!PyUpb_IndexToRange(item, size, &idx, &count, &step)) return NULL; + if (step == 0) { + return PyUpb_GenericSequence_GetItem(_self, idx); + } else { + PyObject* list = PyList_New(count); + for (Py_ssize_t i = 0; i < count; i++, idx += step) { + const void* elem = self->funcs->index(self->parent, idx); + PyList_SetItem(list, i, self->funcs->get_elem_wrapper(elem)); + } + return list; + } +} + // Linear search. Could optimize this in some cases (defs that have index), // but not all (FileDescriptor.dependencies). static int PyUpb_GenericSequence_Find(PyObject* _self, PyObject* item) { @@ -323,10 +344,10 @@ static PyType_Slot PyUpb_GenericSequence_Slots[] = { {Py_sq_length, PyUpb_GenericSequence_Length}, {Py_sq_item, PyUpb_GenericSequence_GetItem}, {Py_tp_richcompare, &PyUpb_GenericSequence_RichCompare}, + {Py_mp_subscript, PyUpb_GenericSequence_Subscript}, // These were implemented for Python/C++ but so far have not been required. // {Py_tp_repr, &PyUpb_GenericSequence_Repr}, // {Py_sq_contains, PyUpb_GenericSequence_Contains}, - // {Py_mp_subscript, PyUpb_GenericSequence_Subscript}, // {Py_mp_ass_subscript, PyUpb_GenericSequence_AssignSubscript}, {0, NULL}, }; diff --git a/upb/python/protobuf.c b/upb/python/protobuf.c index eaba9514ab..9fd858bd85 100644 --- a/upb/python/protobuf.c +++ b/upb/python/protobuf.c @@ -325,6 +325,34 @@ PyObject* PyUpb_Forbidden_New(PyObject* cls, PyObject* args, PyObject* kwds) { return NULL; } +bool PyUpb_IndexToRange(PyObject* index, Py_ssize_t size, Py_ssize_t* i, + Py_ssize_t* count, Py_ssize_t* step) { + assert(i && count && step); + if (PySlice_Check(index)) { + Py_ssize_t start, stop; + if (PySlice_Unpack(index, &start, &stop, step) < 0) return false; + *count = PySlice_AdjustIndices(size, &start, &stop, *step); + *i = start; + } else { + *i = PyNumber_AsSsize_t(index, PyExc_IndexError); + + if (*i == -1 && PyErr_Occurred()) { + PyErr_SetString(PyExc_TypeError, "list indices must be integers"); + return false; + } + + if (*i < 0) *i += size; + *step = 0; + *count = 1; + + if (*i < 0 || size <= *i) { + PyErr_Format(PyExc_IndexError, "list index out of range"); + return false; + } + } + return true; +} + // ----------------------------------------------------------------------------- // Module Entry Point // ----------------------------------------------------------------------------- diff --git a/upb/python/protobuf.h b/upb/python/protobuf.h index 480ff13abe..49f08b17aa 100644 --- a/upb/python/protobuf.h +++ b/upb/python/protobuf.h @@ -231,4 +231,11 @@ static inline PyObject* PyUpb_NewRef(PyObject* obj) { const char* PyUpb_GetStrData(PyObject* obj); const char* PyUpb_VerifyStrData(PyObject* obj); +// For an expression like: +// foo[index] +// +// Converts `index` to an effective i/count/step, for a repeated field +// or descriptor sequence of size 'size'. +bool PyUpb_IndexToRange(PyObject* index, Py_ssize_t size, Py_ssize_t* i, + Py_ssize_t* count, Py_ssize_t* step); #endif // PYUPB_PROTOBUF_H__ diff --git a/upb/python/repeated.c b/upb/python/repeated.c index b526c88875..abb34e880a 100644 --- a/upb/python/repeated.c +++ b/upb/python/repeated.c @@ -39,39 +39,6 @@ static PyObject* PyUpb_RepeatedCompositeContainer_Append(PyObject* _self, static PyObject* PyUpb_RepeatedScalarContainer_Append(PyObject* _self, PyObject* value); -// For an expression like: -// foo[index] -// -// Converts `index` to an effective i/count/step, for a repeated field -// field of size `size`. -static bool IndexToRange(PyObject* index, Py_ssize_t size, Py_ssize_t* i, - Py_ssize_t* count, Py_ssize_t* step) { - assert(i && count && step); - if (PySlice_Check(index)) { - Py_ssize_t start, stop; - if (PySlice_Unpack(index, &start, &stop, step) < 0) return false; - *count = PySlice_AdjustIndices(size, &start, &stop, *step); - *i = start; - } else { - *i = PyNumber_AsSsize_t(index, PyExc_IndexError); - - if (*i == -1 && PyErr_Occurred()) { - PyErr_SetString(PyExc_TypeError, "list indices must be integers"); - return false; - } - - if (*i < 0) *i += size; - *step = 0; - *count = 1; - - if (*i < 0 || size <= *i) { - PyErr_Format(PyExc_IndexError, "list index out of range"); - return false; - } - } - return true; -} - // Wrapper for a repeated field. typedef struct { PyObject_HEAD; @@ -322,7 +289,7 @@ static PyObject* PyUpb_RepeatedContainer_Subscript(PyObject* _self, upb_Array* arr = PyUpb_RepeatedContainer_GetIfReified(self); Py_ssize_t size = arr ? upb_Array_Size(arr) : 0; Py_ssize_t idx, count, step; - if (!IndexToRange(key, size, &idx, &count, &step)) return NULL; + if (!PyUpb_IndexToRange(key, size, &idx, &count, &step)) return NULL; const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self); if (step == 0) { return PyUpb_UpbToPy(upb_Array_Get(arr, idx), f, self->arena); @@ -445,7 +412,7 @@ static int PyUpb_RepeatedContainer_AssignSubscript(PyObject* _self, upb_Array* arr = PyUpb_RepeatedContainer_EnsureReified(_self); Py_ssize_t size = arr ? upb_Array_Size(arr) : 0; Py_ssize_t idx, count, step; - if (!IndexToRange(key, size, &idx, &count, &step)) return -1; + if (!PyUpb_IndexToRange(key, size, &idx, &count, &step)) return -1; if (value) { return PyUpb_RepeatedContainer_SetSubscript(self, arr, f, idx, count, step, value);