Add upb python Py_mp_subscript for descriptor sequences.

A behavior diff that between upb and the other two implementations.

PiperOrigin-RevId: 561471616

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."""

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},
 };

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
 // -----------------------------------------------------------------------------

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__

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);