diff --git a/Lib/test/test_capi/test_misc.py b/Lib/test/test_capi/test_misc.py index 9a5a1d938e2869f..b9c6882444ab767 100644 --- a/Lib/test/test_capi/test_misc.py +++ b/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 diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c index 033ee3234cf604d..00ba09b6e2fe248 100644 --- a/Modules/_testcapimodule.c +++ b/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},