gh-94808: add tests covering PySequence_[InPlace_]Repeat

Lib/test/test_capi/test_misc.py

@@ -404,6 +404,119 @@ def test_buildvalue_ints(self):
     def test_buildvalue_N(self):
         _testcapi.test_buildvalue_N()
 
+    def test_sequence_repeat_subtypes(self):
+        # List subtype:
+        class CustomList(list): ...
+
+        sequence = CustomList(['a', 'b'])
+        repeated_sequence = _testcapi.sequence_repeat(sequence, 3)
+        self.assertIs(type(repeated_sequence), list)
+        self.assertEqual(repeated_sequence, sequence * 3)
+        self.assertEqual(sequence, CustomList(['a', 'b']))
+
+        # Tuple subtype:
+        class CustomTuple(tuple): ...
+
+        sequence = CustomTuple(('a', 'b'))
+        repeated_sequence = _testcapi.sequence_repeat(sequence, 3)
+        self.assertIs(type(repeated_sequence), tuple)
+        self.assertEqual(repeated_sequence, sequence * 3)
+        self.assertEqual(sequence, CustomTuple(('a', 'b')))
+
+        # String subtype:
+        class CustomStr(str): ...
+
+        sequence = CustomStr('abc')
+        repeated_sequence = _testcapi.sequence_repeat(sequence, 2)
+        self.assertIs(type(repeated_sequence), str)
+        self.assertEqual(repeated_sequence, sequence * 2)
+        self.assertEqual(sequence, CustomStr('abc'))
+
+    def test_sequence_repeat_numeric_fallback(self):
+        # PySequence_Repeat falls back to nb_multiply (__mul__) when the
+        # object looks like a sequence but has no sq_repeat slot.
+        class CustomSequenceWithMul:
+            def __getitem__(self, index):
+                return index  # so that it looks like a sequence
+
+            def __mul__(self, count):
+                return f'repeated {count} times'
+
+        sequence = CustomSequenceWithMul()
+        self.assertEqual(_testcapi.sequence_repeat(sequence, 3),
+                         'repeated 3 times')
+
+    def test_sequence_repeat_errors(self):
+        # Corner case, custom subtype returns `NotImplemented`:
+        class CustomListWithNI(list):
+            def __mul__(self, other):
+                return NotImplemented
+
+        sequence = CustomListWithNI(['a', 'b'])
+        with self.assertRaises(TypeError):
+            _testcapi.sequence_repeat(sequence, 3)
+        self.assertEqual(sequence, CustomListWithNI(['a', 'b']))  # not changed
+
+        # Non-sequences do not work (even with `__mul__`):
+        class CustomTypeWithMul:
+            def __mul__(self, other):
+                return 5  # won't be called
+
+        not_sequence = CustomTypeWithMul()
+        with self.assertRaises(TypeError):
+            _testcapi.sequence_repeat(not_sequence, 1)
+
+    def test_sequence_inplace_repeat_subtypes(self):
+        # List subtype:
+        class CustomList(list): ...
+
+        sequence = CustomList(['a', 'b'])
+        repeated_sequence = _testcapi.sequence_inplace_repeat(sequence, 3)
+        self.assertIs(repeated_sequence, sequence)
+        self.assertEqual(repeated_sequence, ['a', 'b', 'a', 'b', 'a', 'b'])
+
+        # Tuple subtype:
+        class CustomTuple(tuple): ...
+
+        sequence = CustomTuple(('a', 'b'))
+        repeated_sequence = _testcapi.sequence_inplace_repeat(sequence, 3)
+        self.assertIs(type(repeated_sequence), tuple)
+        self.assertEqual(repeated_sequence, sequence * 3)
+        self.assertEqual(sequence, CustomTuple(('a', 'b')))
+
+        # String subtype:
+        class CustomStr(str): ...
+
+        sequence = CustomStr('abc')
+        repeated_sequence = _testcapi.sequence_inplace_repeat(sequence, 2)
+        self.assertIs(type(repeated_sequence), str)
+        self.assertEqual(repeated_sequence, sequence * 2)
+        self.assertEqual(sequence, CustomStr('abc'))
+
+    def test_sequence_inplace_repeat_errors(self):
+        # Corner case, custom subtype returns `NotImplemented`:
+        class CustomListWithNI(list):
+            def __mul__(self, other):
+                return NotImplemented
+            def __imul__(self, other):
+                return NotImplemented
+
+        sequence = CustomListWithNI(['a', 'b'])
+        with self.assertRaises(TypeError):
+            _testcapi.sequence_inplace_repeat(sequence, 3)
+        self.assertEqual(sequence, CustomListWithNI(['a', 'b']))  # not changed
+
+        # Non-sequences do not work (even with `__mul__` and `__imul__`):
+        class CustomTypeWithMul:
+            def __mul__(self, other):
+                return 0
+            def __imul__(self, other):
+                return 1  # won't be called
+
+        not_sequence = CustomTypeWithMul()
+        with self.assertRaises(TypeError):
+            _testcapi.sequence_inplace_repeat(not_sequence, 1)
+
     def test_trashcan_subclass(self):
         # bpo-35983: Check that the trashcan mechanism for "list" is NOT
         # activated when its tp_dealloc is being called by a subclass

Modules/_testcapimodule.c

@@ -1491,6 +1491,30 @@ py_w_stopcode(PyObject *self, PyObject *args)
 #endif
 
 
+static PyObject *
+sequence_repeat(PyObject* self, PyObject *args)
+{
+    PyObject *sequence;
+    Py_ssize_t count;
+    if (!PyArg_ParseTuple(args, "On", &sequence, &count)) {
+        return NULL;
+    }
+
+    return PySequence_Repeat(sequence, count);
+}
+
+static PyObject *
+sequence_inplace_repeat(PyObject* self, PyObject *args)
+{
+    PyObject *sequence;
+    Py_ssize_t count;
+    if (!PyArg_ParseTuple(args, "On", &sequence, &count)) {
+        return NULL;
+    }
+
+    return PySequence_InPlaceRepeat(sequence, count);
+}
+
 static PyObject *
 test_pythread_tss_key_state(PyObject *self, PyObject *args)
 {
@@ -2949,6 +2973,8 @@ static PyMethodDef TestMethods[] = {
 #ifdef W_STOPCODE
     {"W_STOPCODE", py_w_stopcode, METH_VARARGS},
 #endif
+    {"sequence_repeat", sequence_repeat, METH_VARARGS},
+    {"sequence_inplace_repeat", sequence_inplace_repeat, METH_VARARGS},
     {"test_pythread_tss_key_state", test_pythread_tss_key_state, METH_VARARGS},
     {"bad_get", bad_get, METH_VARARGS},
     {"meth_varargs", meth_varargs, METH_VARARGS},