mirror of https://github.com/python/cpython
Port r69837: Fix keyword arguments for itertools.count(). Step arg without a start arg was ignored.
This commit is contained in:
parent
bd171bcfc4
commit
eb13fdda59
|
@ -354,6 +354,10 @@ class TestBasicOps(unittest.TestCase):
|
||||||
|
|
||||||
def test_count_with_stride(self):
|
def test_count_with_stride(self):
|
||||||
self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
|
self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
|
||||||
|
self.assertEqual(lzip('abc',count(start=2,step=3)),
|
||||||
|
[('a', 2), ('b', 5), ('c', 8)])
|
||||||
|
self.assertEqual(lzip('abc',count(step=-1)),
|
||||||
|
[('a', 0), ('b', -1), ('c', -2)])
|
||||||
self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
|
self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
|
||||||
self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
|
self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
|
||||||
self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
|
self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
|
||||||
|
|
|
@ -2892,19 +2892,19 @@ typedef struct {
|
||||||
|
|
||||||
/* Counting logic and invariants:
|
/* Counting logic and invariants:
|
||||||
|
|
||||||
C_add_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
|
fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
|
||||||
|
|
||||||
assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
|
assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
|
||||||
Advances with: cnt += 1
|
Advances with: cnt += 1
|
||||||
When count hits Y_SSIZE_T_MAX, switch to Py_add_mode.
|
When count hits Y_SSIZE_T_MAX, switch to slow_mode.
|
||||||
|
|
||||||
Py_add_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
|
slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
|
||||||
|
|
||||||
assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
|
assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
|
||||||
All counting is done with python objects (no overflows or underflows).
|
All counting is done with python objects (no overflows or underflows).
|
||||||
Advances with: long_cnt += long_step
|
Advances with: long_cnt += long_step
|
||||||
Step may be zero -- effectively a slow version of repeat(cnt).
|
Step may be zero -- effectively a slow version of repeat(cnt).
|
||||||
Either long_cnt or long_step may be a float.
|
Either long_cnt or long_step may be a float, Fraction, or Decimal.
|
||||||
*/
|
*/
|
||||||
|
|
||||||
static PyTypeObject count_type;
|
static PyTypeObject count_type;
|
||||||
|
@ -2913,6 +2913,7 @@ static PyObject *
|
||||||
count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
|
count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
|
||||||
{
|
{
|
||||||
countobject *lz;
|
countobject *lz;
|
||||||
|
int slow_mode = 0;
|
||||||
Py_ssize_t cnt = 0;
|
Py_ssize_t cnt = 0;
|
||||||
PyObject *long_cnt = NULL;
|
PyObject *long_cnt = NULL;
|
||||||
PyObject *long_step = NULL;
|
PyObject *long_step = NULL;
|
||||||
|
@ -2922,36 +2923,51 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
|
||||||
kwlist, &long_cnt, &long_step))
|
kwlist, &long_cnt, &long_step))
|
||||||
return NULL;
|
return NULL;
|
||||||
|
|
||||||
if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
|
if (long_cnt != NULL && !PyNumber_Check(long_cnt) ||
|
||||||
(long_step != NULL && !PyNumber_Check(long_step))) {
|
long_step != NULL && !PyNumber_Check(long_step)) {
|
||||||
PyErr_SetString(PyExc_TypeError, "a number is required");
|
PyErr_SetString(PyExc_TypeError, "a number is required");
|
||||||
return NULL;
|
return NULL;
|
||||||
}
|
}
|
||||||
|
|
||||||
if (long_step == NULL) {
|
|
||||||
/* If not specified, step defaults to 1 */
|
|
||||||
long_step = PyLong_FromLong(1);
|
|
||||||
if (long_step == NULL)
|
|
||||||
return NULL;
|
|
||||||
} else
|
|
||||||
Py_INCREF(long_step);
|
|
||||||
assert(long_step != NULL);
|
|
||||||
|
|
||||||
if (long_cnt != NULL) {
|
if (long_cnt != NULL) {
|
||||||
cnt = PyLong_AsSsize_t(long_cnt);
|
cnt = PyLong_AsSsize_t(long_cnt);
|
||||||
if ((cnt == -1 && PyErr_Occurred()) ||
|
if (cnt == -1 && PyErr_Occurred() || !PyLong_Check(long_cnt)) {
|
||||||
!PyIndex_Check(long_cnt) ||
|
|
||||||
!PyLong_Check(long_step) ||
|
|
||||||
PyLong_AS_LONG(long_step) != 1) {
|
|
||||||
/* Switch to Py_add_mode */
|
|
||||||
PyErr_Clear();
|
PyErr_Clear();
|
||||||
Py_INCREF(long_cnt);
|
slow_mode = 1;
|
||||||
cnt = PY_SSIZE_T_MAX;
|
|
||||||
} else
|
|
||||||
long_cnt = NULL;
|
|
||||||
}
|
}
|
||||||
assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL) ||
|
Py_INCREF(long_cnt);
|
||||||
(cnt == PY_SSIZE_T_MAX && long_cnt != NULL));
|
} else {
|
||||||
|
cnt = 0;
|
||||||
|
long_cnt = PyLong_FromLong(0);
|
||||||
|
}
|
||||||
|
|
||||||
|
/* If not specified, step defaults to 1 */
|
||||||
|
if (long_step == NULL) {
|
||||||
|
long_step = PyLong_FromLong(1);
|
||||||
|
if (long_step == NULL) {
|
||||||
|
Py_DECREF(long_cnt);
|
||||||
|
return NULL;
|
||||||
|
}
|
||||||
|
} else
|
||||||
|
Py_INCREF(long_step);
|
||||||
|
|
||||||
|
assert(long_cnt != NULL && long_step != NULL);
|
||||||
|
|
||||||
|
/* Fast mode only works when the step is 1 */
|
||||||
|
if (!PyLong_Check(long_step) ||
|
||||||
|
PyLong_AS_LONG(long_step) != 1) {
|
||||||
|
slow_mode = 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (slow_mode)
|
||||||
|
cnt = PY_SSIZE_T_MAX;
|
||||||
|
else
|
||||||
|
Py_CLEAR(long_cnt);
|
||||||
|
|
||||||
|
assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode ||
|
||||||
|
cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode);
|
||||||
|
assert(slow_mode ||
|
||||||
|
PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1);
|
||||||
|
|
||||||
/* create countobject structure */
|
/* create countobject structure */
|
||||||
lz = (countobject *)type->tp_alloc(type, 0);
|
lz = (countobject *)type->tp_alloc(type, 0);
|
||||||
|
@ -2991,7 +3007,7 @@ count_nextlong(countobject *lz)
|
||||||
|
|
||||||
long_cnt = lz->long_cnt;
|
long_cnt = lz->long_cnt;
|
||||||
if (long_cnt == NULL) {
|
if (long_cnt == NULL) {
|
||||||
/* Switch to Py_add_mode */
|
/* Switch to slow_mode */
|
||||||
long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
|
long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
|
||||||
if (long_cnt == NULL)
|
if (long_cnt == NULL)
|
||||||
return NULL;
|
return NULL;
|
||||||
|
@ -3034,11 +3050,10 @@ count_repr(countobject *lz)
|
||||||
}
|
}
|
||||||
|
|
||||||
PyDoc_STRVAR(count_doc,
|
PyDoc_STRVAR(count_doc,
|
||||||
"count([start[, step]]) --> count object\n\
|
"count(start=0, step=1]) --> count object\n\
|
||||||
\n\
|
\n\
|
||||||
Return a count object whose .__next__() method returns consecutive\n\
|
Return a count object whose .__next__() method returns consecutive values.\n\
|
||||||
integers starting from zero or, if specified, from start.\n\
|
Equivalent to:\n\n\
|
||||||
If step is specified, counts by that interval. Equivalent to:\n\n\
|
|
||||||
def count(firstval=0, step=1):\n\
|
def count(firstval=0, step=1):\n\
|
||||||
x = firstval\n\
|
x = firstval\n\
|
||||||
while 1:\n\
|
while 1:\n\
|
||||||
|
|
Loading…
Reference in New Issue