From a2e280e86535a9c46184eccdb9e6295ad156ee69 Mon Sep 17 00:00:00 2001
From: Joshua Haberman <jhaberman@gmail.com>
Date: Sun, 9 Jan 2022 12:03:17 -0800
Subject: [PATCH] Fixed slice assignment of unequal length.

---
 .../pb_unit_tests/reflection_test_wrapper.py  |  1 -
 python/repeated.c                             | 21 +++++++++++++------
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/python/pb_unit_tests/reflection_test_wrapper.py b/python/pb_unit_tests/reflection_test_wrapper.py
index 7969abaeaa..8f7d48f47b 100644
--- a/python/pb_unit_tests/reflection_test_wrapper.py
+++ b/python/pb_unit_tests/reflection_test_wrapper.py
@@ -35,7 +35,6 @@ reflection_test.Proto2ReflectionTest.testListFieldsAndExtensions.__unittest_expe
 reflection_test.Proto2ReflectionTest.testRepeatedCompositeReverse_Empty.__unittest_expecting_failure__ = True
 reflection_test.Proto2ReflectionTest.testRepeatedCompositeReverse_NonEmpty.__unittest_expecting_failure__ = True
 reflection_test.Proto2ReflectionTest.testRepeatedListExtensions.__unittest_expecting_failure__ = True
-reflection_test.Proto2ReflectionTest.testRepeatedScalars.__unittest_expecting_failure__ = True
 reflection_test.Proto2ReflectionTest.testRepeatedScalarsReverse_Empty.__unittest_expecting_failure__ = True
 reflection_test.Proto2ReflectionTest.testRepeatedScalarsReverse_NonEmpty.__unittest_expecting_failure__ = True
 reflection_test.Proto2ReflectionTest.testSingularListExtensions.__unittest_expecting_failure__ = True
diff --git a/python/repeated.c b/python/repeated.c
index 1b32fd4f8f..b758bef402 100644
--- a/python/repeated.c
+++ b/python/repeated.c
@@ -359,12 +359,21 @@ static int PyUpb_RepeatedContainer_SetSubscript(
   PyObject* item = NULL;
   int ret = -1;
   if (!seq) goto err;
-  if (PySequence_Size(seq) != count) {
-    PyErr_Format(PyExc_ValueError,
-                  "attempt to assign sequence of size %zd to extended slice "
-                  "of size %zd",
-                  PySequence_Size(seq), count);
-    goto err;
+  Py_ssize_t seq_size = PySequence_Size(seq);
+  if (seq_size != count) {
+    if (step == 1) {
+      // We must shift the tail elements (either right or left).
+      size_t tail = upb_array_size(arr) - (idx + count);
+      upb_array_resize(arr, idx + seq_size + tail, arena);
+      upb_array_move(arr, idx + seq_size, idx + count, tail);
+      count = seq_size;
+    } else {
+      PyErr_Format(PyExc_ValueError,
+                   "attempt to assign sequence of  %zd to extended slice "
+                   "of size %zd",
+                   seq_size, count);
+      goto err;
+    }
   }
   for (Py_ssize_t i = 0; i < count; i++, idx += step) {
     upb_msgval msgval;