blob: 66d743fb9a5c1677c3d4688dc61e265f42bdd5f6 [file] [log] [blame]
Guido van Rossum12d12c51993-10-26 17:58:25 +00001/* Range object implementation */
2
Guido van Rossumc0b618a1997-05-02 03:12:38 +00003#include "Python.h"
Thomas Woutersefafcea2001-07-09 12:30:54 +00004
Guido van Rossum805365e2007-05-07 22:24:25 +00005/* Support objects whose length is > PY_SSIZE_T_MAX.
6
7 This could be sped up for small PyLongs if they fit in an Py_ssize_t.
8 This only matters on Win64. Though we could use PY_LONG_LONG which
9 would presumably help perf.
10*/
11
Guido van Rossum12d12c51993-10-26 17:58:25 +000012typedef struct {
Guido van Rossum805365e2007-05-07 22:24:25 +000013 PyObject_HEAD
14 PyObject *start;
15 PyObject *stop;
16 PyObject *step;
Guido van Rossum12d12c51993-10-26 17:58:25 +000017} rangeobject;
18
Guido van Rossum805365e2007-05-07 22:24:25 +000019/* Helper function for validating step. Always returns a new reference or
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000020 NULL on error.
Guido van Rossum805365e2007-05-07 22:24:25 +000021*/
22static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +000023validate_step(PyObject *step) {
Guido van Rossum805365e2007-05-07 22:24:25 +000024 /* No step specified, use a step of 1. */
25 if (!step)
Christian Heimes217cfd12007-12-02 14:31:20 +000026 return PyLong_FromLong(1);
Guido van Rossum805365e2007-05-07 22:24:25 +000027
28 step = PyNumber_Index(step);
29 if (step) {
30 Py_ssize_t istep = PyNumber_AsSsize_t(step, NULL);
31 if (istep == -1 && PyErr_Occurred()) {
32 /* Ignore OverflowError, we know the value isn't 0. */
33 PyErr_Clear();
34 }
35 else if (istep == 0) {
36 PyErr_SetString(PyExc_ValueError,
37 "range() arg 3 must not be zero");
38 Py_CLEAR(step);
39 }
40 }
41
42 return step;
Raymond Hettingerc4c453f2002-06-05 23:12:45 +000043}
44
Guido van Rossum805365e2007-05-07 22:24:25 +000045/* XXX(nnorwitz): should we error check if the user passes any empty ranges?
46 range(-10)
47 range(0, -5)
48 range(0, 5, -1)
49*/
Raymond Hettingerc4c453f2002-06-05 23:12:45 +000050static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +000051range_new(PyTypeObject *type, PyObject *args, PyObject *kw) {
Guido van Rossum805365e2007-05-07 22:24:25 +000052 rangeobject *obj = NULL;
53 PyObject *start = NULL, *stop = NULL, *step = NULL;
Raymond Hettingerc4c453f2002-06-05 23:12:45 +000054
Guido van Rossum805365e2007-05-07 22:24:25 +000055 if (!_PyArg_NoKeywords("range()", kw))
56 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +000057
Guido van Rossum805365e2007-05-07 22:24:25 +000058 if (PyTuple_Size(args) <= 1) {
59 if (!PyArg_UnpackTuple(args, "range", 1, 1, &stop))
Raymond Hettinger94f55832009-06-12 18:40:16 +000060 return NULL;
Guido van Rossum805365e2007-05-07 22:24:25 +000061 stop = PyNumber_Index(stop);
62 if (!stop)
Raymond Hettinger94f55832009-06-12 18:40:16 +000063 return NULL;
Christian Heimes217cfd12007-12-02 14:31:20 +000064 start = PyLong_FromLong(0);
Raymond Hettinger94f55832009-06-12 18:40:16 +000065 if (!start) {
66 Py_DECREF(stop);
67 return NULL;
68 }
Christian Heimes217cfd12007-12-02 14:31:20 +000069 step = PyLong_FromLong(1);
Raymond Hettinger94f55832009-06-12 18:40:16 +000070 if (!step) {
71 Py_DECREF(stop);
72 Py_DECREF(start);
73 return NULL;
74 }
Guido van Rossum805365e2007-05-07 22:24:25 +000075 }
76 else {
77 if (!PyArg_UnpackTuple(args, "range", 2, 3,
78 &start, &stop, &step))
Raymond Hettinger94f55832009-06-12 18:40:16 +000079 return NULL;
Tim Petersfeec4532004-08-08 07:17:39 +000080
Guido van Rossum805365e2007-05-07 22:24:25 +000081 /* Convert borrowed refs to owned refs */
82 start = PyNumber_Index(start);
Raymond Hettinger94f55832009-06-12 18:40:16 +000083 if (!start)
84 return NULL;
Guido van Rossum805365e2007-05-07 22:24:25 +000085 stop = PyNumber_Index(stop);
Raymond Hettinger94f55832009-06-12 18:40:16 +000086 if (!stop) {
87 Py_DECREF(start);
88 return NULL;
89 }
90 step = validate_step(step); /* Caution, this can clear exceptions */
91 if (!step) {
92 Py_DECREF(start);
93 Py_DECREF(stop);
94 return NULL;
95 }
Guido van Rossum805365e2007-05-07 22:24:25 +000096 }
97
98 obj = PyObject_New(rangeobject, &PyRange_Type);
99 if (obj == NULL)
100 goto Fail;
101 obj->start = start;
102 obj->stop = stop;
103 obj->step = step;
104 return (PyObject *) obj;
105
106Fail:
107 Py_XDECREF(start);
108 Py_XDECREF(stop);
109 Py_XDECREF(step);
110 return NULL;
Raymond Hettingerc4c453f2002-06-05 23:12:45 +0000111}
112
Martin v. Löwis14f8b4c2002-06-13 20:33:02 +0000113PyDoc_STRVAR(range_doc,
Guido van Rossum805365e2007-05-07 22:24:25 +0000114"range([start,] stop[, step]) -> range object\n\
Raymond Hettingerc4c453f2002-06-05 23:12:45 +0000115\n\
Guido van Rossum805365e2007-05-07 22:24:25 +0000116Returns an iterator that generates the numbers in the range on demand.");
Raymond Hettingerc4c453f2002-06-05 23:12:45 +0000117
Guido van Rossum805365e2007-05-07 22:24:25 +0000118static void
Benjamin Peterson155614b2010-11-20 22:44:32 +0000119range_dealloc(rangeobject *r) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000120 Py_DECREF(r->start);
121 Py_DECREF(r->stop);
122 Py_DECREF(r->step);
Amaury Forgeot d'Arcb7f17e42007-11-15 20:52:21 +0000123 PyObject_Del(r);
Guido van Rossum805365e2007-05-07 22:24:25 +0000124}
125
Mark Dickinson732166d2009-06-28 22:08:40 +0000126/* Return number of items in range (lo, hi, step) as a PyLong object,
127 * when arguments are PyLong objects. Arguments MUST return 1 with
128 * PyLong_Check(). Return NULL when there is an error.
Guido van Rossum805365e2007-05-07 22:24:25 +0000129 */
130static PyObject*
Benjamin Peterson155614b2010-11-20 22:44:32 +0000131range_length_obj(rangeobject *r) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000132 /* -------------------------------------------------------------
133 Algorithm is equal to that of get_len_of_range(), but it operates
Benjamin Petersona47af9c2009-06-24 21:14:38 +0000134 on PyObjects (which are assumed to be PyLong objects).
Guido van Rossum805365e2007-05-07 22:24:25 +0000135 ---------------------------------------------------------------*/
Mark Dickinson211c6252009-02-01 10:28:51 +0000136 int cmp_result;
Guido van Rossum805365e2007-05-07 22:24:25 +0000137 PyObject *lo, *hi;
138 PyObject *step = NULL;
139 PyObject *diff = NULL;
140 PyObject *one = NULL;
141 PyObject *tmp1 = NULL, *tmp2 = NULL, *result;
142 /* holds sub-expression evaluations */
143
144 PyObject *zero = PyLong_FromLong(0);
145 if (zero == NULL)
146 return NULL;
Mark Dickinson211c6252009-02-01 10:28:51 +0000147 cmp_result = PyObject_RichCompareBool(r->step, zero, Py_GT);
Guido van Rossum805365e2007-05-07 22:24:25 +0000148 Py_DECREF(zero);
Mark Dickinson211c6252009-02-01 10:28:51 +0000149 if (cmp_result == -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000150 return NULL;
151
Mark Dickinson211c6252009-02-01 10:28:51 +0000152 if (cmp_result == 1) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000153 lo = r->start;
154 hi = r->stop;
155 step = r->step;
156 Py_INCREF(step);
157 } else {
158 lo = r->stop;
159 hi = r->start;
160 step = PyNumber_Negative(r->step);
161 if (!step)
162 return NULL;
163 }
164
165 /* if (lo >= hi), return length of 0. */
Mark Dickinson211c6252009-02-01 10:28:51 +0000166 if (PyObject_RichCompareBool(lo, hi, Py_GE) == 1) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000167 Py_XDECREF(step);
168 return PyLong_FromLong(0);
169 }
170
171 if ((one = PyLong_FromLong(1L)) == NULL)
172 goto Fail;
173
174 if ((tmp1 = PyNumber_Subtract(hi, lo)) == NULL)
175 goto Fail;
176
177 if ((diff = PyNumber_Subtract(tmp1, one)) == NULL)
178 goto Fail;
179
180 if ((tmp2 = PyNumber_FloorDivide(diff, step)) == NULL)
181 goto Fail;
182
183 if ((result = PyNumber_Add(tmp2, one)) == NULL)
184 goto Fail;
185
186 Py_DECREF(tmp2);
187 Py_DECREF(diff);
188 Py_DECREF(step);
189 Py_DECREF(tmp1);
190 Py_DECREF(one);
191 return result;
192
193 Fail:
194 Py_XDECREF(tmp2);
195 Py_XDECREF(diff);
196 Py_XDECREF(step);
197 Py_XDECREF(tmp1);
198 Py_XDECREF(one);
199 return NULL;
Guido van Rossum12d12c51993-10-26 17:58:25 +0000200}
201
Martin v. Löwis18e16552006-02-15 17:27:45 +0000202static Py_ssize_t
Benjamin Peterson155614b2010-11-20 22:44:32 +0000203range_length(rangeobject *r) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000204 PyObject *len = range_length_obj(r);
205 Py_ssize_t result = -1;
206 if (len) {
207 result = PyLong_AsSsize_t(len);
208 Py_DECREF(len);
209 }
210 return result;
211}
212
213/* range(...)[x] is necessary for: seq[:] = range(...) */
214
215static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000216range_item(rangeobject *r, Py_ssize_t i) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000217 Py_ssize_t len = range_length(r);
218 PyObject *rem, *incr, *result;
219
220 /* XXX(nnorwitz): should negative indices be supported? */
221 /* XXX(nnorwitz): should support range[x] where x > PY_SSIZE_T_MAX? */
222 if (i < 0 || i >= len) {
223 if (!PyErr_Occurred())
224 PyErr_SetString(PyExc_IndexError,
Walter Dörwald4ad94212007-05-21 18:01:17 +0000225 "range object index out of range");
Benjamin Petersondf0a5cb2008-04-25 21:15:37 +0000226 return NULL;
227 }
Guido van Rossum805365e2007-05-07 22:24:25 +0000228
229 /* XXX(nnorwitz): optimize for short ints. */
Raymond Hettingerad3f3322008-02-09 04:13:49 +0000230 rem = PyLong_FromSsize_t(i);
Guido van Rossum805365e2007-05-07 22:24:25 +0000231 if (!rem)
232 return NULL;
233 incr = PyNumber_Multiply(rem, r->step);
234 Py_DECREF(rem);
235 if (!incr)
236 return NULL;
237 result = PyNumber_Add(r->start, incr);
238 Py_DECREF(incr);
239 return result;
Guido van Rossum7d6aa511993-12-21 22:50:31 +0000240}
241
Guido van Rossumc0b618a1997-05-02 03:12:38 +0000242static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000243range_repr(rangeobject *r) {
Walter Dörwald03b43d82007-05-21 10:43:34 +0000244 Py_ssize_t istep;
Tim Petersd976ab72004-08-08 06:29:10 +0000245
Guido van Rossum805365e2007-05-07 22:24:25 +0000246 /* Check for special case values for printing. We don't always
Walter Dörwald03b43d82007-05-21 10:43:34 +0000247 need the step value. We don't care about errors
Guido van Rossum805365e2007-05-07 22:24:25 +0000248 (it means overflow), so clear the errors. */
Guido van Rossum805365e2007-05-07 22:24:25 +0000249 istep = PyNumber_AsSsize_t(r->step, NULL);
250 if (istep != 1 || (istep == -1 && PyErr_Occurred())) {
251 PyErr_Clear();
Guido van Rossum805365e2007-05-07 22:24:25 +0000252 }
253
Walter Dörwald03b43d82007-05-21 10:43:34 +0000254 if (istep == 1)
Walter Dörwald850e5162007-05-20 08:19:54 +0000255 return PyUnicode_FromFormat("range(%R, %R)", r->start, r->stop);
256 else
257 return PyUnicode_FromFormat("range(%R, %R, %R)",
258 r->start, r->stop, r->step);
Thomas Woutersefafcea2001-07-09 12:30:54 +0000259}
260
Alexandre Vassalotti75056072008-06-10 04:03:04 +0000261/* Pickling support */
262static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000263range_reduce(rangeobject *r, PyObject *args) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000264 return Py_BuildValue("(O(OOO))", Py_TYPE(r),
Alexandre Vassalotti75056072008-06-10 04:03:04 +0000265 r->start, r->stop, r->step);
266}
267
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000268/* Assumes (PyLong_CheckExact(ob) || PyBool_Check(ob)) */
269static int
Benjamin Peterson155614b2010-11-20 22:44:32 +0000270range_contains_long(rangeobject *r, PyObject *ob) {
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000271 int cmp1, cmp2, cmp3;
272 PyObject *tmp1 = NULL;
273 PyObject *tmp2 = NULL;
274 PyObject *zero = NULL;
275 int result = -1;
276
277 zero = PyLong_FromLong(0);
278 if (zero == NULL) /* MemoryError in int(0) */
279 goto end;
280
281 /* Check if the value can possibly be in the range. */
282
283 cmp1 = PyObject_RichCompareBool(r->step, zero, Py_GT);
284 if (cmp1 == -1)
285 goto end;
286 if (cmp1 == 1) { /* positive steps: start <= ob < stop */
287 cmp2 = PyObject_RichCompareBool(r->start, ob, Py_LE);
288 cmp3 = PyObject_RichCompareBool(ob, r->stop, Py_LT);
289 }
290 else { /* negative steps: stop < ob <= start */
291 cmp2 = PyObject_RichCompareBool(ob, r->start, Py_LE);
292 cmp3 = PyObject_RichCompareBool(r->stop, ob, Py_LT);
293 }
294
295 if (cmp2 == -1 || cmp3 == -1) /* TypeError */
296 goto end;
297 if (cmp2 == 0 || cmp3 == 0) { /* ob outside of range */
298 result = 0;
299 goto end;
300 }
301
302 /* Check that the stride does not invalidate ob's membership. */
303 tmp1 = PyNumber_Subtract(ob, r->start);
304 if (tmp1 == NULL)
305 goto end;
306 tmp2 = PyNumber_Remainder(tmp1, r->step);
307 if (tmp2 == NULL)
308 goto end;
309 /* result = (int(ob) - start % step) == 0 */
310 result = PyObject_RichCompareBool(tmp2, zero, Py_EQ);
311 end:
312 Py_XDECREF(tmp1);
313 Py_XDECREF(tmp2);
314 Py_XDECREF(zero);
315 return result;
316}
317
Mark Dickinson3e124ae2009-09-22 21:47:24 +0000318static int
319range_contains(rangeobject *r, PyObject *ob) {
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000320 if (PyLong_CheckExact(ob) || PyBool_Check(ob))
321 return range_contains_long(r, ob);
Mark Dickinson3e124ae2009-09-22 21:47:24 +0000322
Mark Dickinson3e124ae2009-09-22 21:47:24 +0000323 return (int)_PySequence_IterSearch((PyObject*)r, ob,
324 PY_ITERSEARCH_CONTAINS);
325}
326
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000327static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000328range_count(rangeobject *r, PyObject *ob) {
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000329 if (PyLong_CheckExact(ob) || PyBool_Check(ob)) {
Georg Brandl7e5343b2010-11-20 22:40:10 +0000330 int result = range_contains_long(r, ob);
331 if (result == -1)
332 return NULL;
333 else if (result)
Benjamin Peterson0b458d52010-11-20 22:35:41 +0000334 return PyLong_FromLong(1);
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000335 else
Benjamin Peterson0b458d52010-11-20 22:35:41 +0000336 return PyLong_FromLong(0);
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000337 } else {
338 Py_ssize_t count;
339 count = _PySequence_IterSearch((PyObject*)r, ob, PY_ITERSEARCH_COUNT);
340 if (count == -1)
341 return NULL;
342 return PyLong_FromSsize_t(count);
343 }
344}
345
346static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000347range_index(rangeobject *r, PyObject *ob) {
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000348 int contains;
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000349
350 if (!PyLong_CheckExact(ob) && !PyBool_Check(ob)) {
351 Py_ssize_t index;
352 index = _PySequence_IterSearch((PyObject*)r, ob, PY_ITERSEARCH_INDEX);
353 if (index == -1)
354 return NULL;
355 return PyLong_FromSsize_t(index);
356 }
357
358 contains = range_contains_long(r, ob);
359 if (contains == -1)
360 return NULL;
361
Benjamin Peterson155614b2010-11-20 22:44:32 +0000362 if (contains) {
363 PyObject *idx, *tmp = PyNumber_Subtract(ob, r->start);
364 if (tmp == NULL)
365 return NULL;
366 /* idx = (ob - r.start) // r.step */
367 idx = PyNumber_FloorDivide(tmp, r->step);
368 Py_DECREF(tmp);
369 return idx;
370 }
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000371
372 /* object is not in the range */
Benjamin Peterson155614b2010-11-20 22:44:32 +0000373 PyErr_Format(PyExc_ValueError, "%R is not in range", ob);
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000374 return NULL;
375}
376
Guido van Rossumc0b618a1997-05-02 03:12:38 +0000377static PySequenceMethods range_as_sequence = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000378 (lenfunc)range_length, /* sq_length */
379 0, /* sq_concat */
380 0, /* sq_repeat */
Guido van Rossum805365e2007-05-07 22:24:25 +0000381 (ssizeargfunc)range_item, /* sq_item */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000382 0, /* sq_slice */
Mark Dickinson3e124ae2009-09-22 21:47:24 +0000383 0, /* sq_ass_item */
384 0, /* sq_ass_slice */
385 (objobjproc)range_contains, /* sq_contains */
Guido van Rossum12d12c51993-10-26 17:58:25 +0000386};
387
Jeremy Hylton938ace62002-07-17 16:30:39 +0000388static PyObject * range_iter(PyObject *seq);
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000389static PyObject * range_reverse(PyObject *seq);
390
391PyDoc_STRVAR(reverse_doc,
392"Returns a reverse iterator.");
393
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000394PyDoc_STRVAR(count_doc,
395"rangeobject.count(value) -> integer -- return number of occurrences of value");
396
397PyDoc_STRVAR(index_doc,
398"rangeobject.index(value, [start, [stop]]) -> integer -- return index of value.\n"
399"Raises ValueError if the value is not present.");
400
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000401static PyMethodDef range_methods[] = {
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000402 {"__reversed__", (PyCFunction)range_reverse, METH_NOARGS, reverse_doc},
403 {"__reduce__", (PyCFunction)range_reduce, METH_VARARGS},
404 {"count", (PyCFunction)range_count, METH_O, count_doc},
405 {"index", (PyCFunction)range_index, METH_O, index_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000406 {NULL, NULL} /* sentinel */
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000407};
Raymond Hettinger48165d42002-06-05 20:08:48 +0000408
Guido van Rossumc0b618a1997-05-02 03:12:38 +0000409PyTypeObject PyRange_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000410 PyVarObject_HEAD_INIT(&PyType_Type, 0)
411 "range", /* Name of this type */
412 sizeof(rangeobject), /* Basic object size */
413 0, /* Item size for varobject */
414 (destructor)range_dealloc, /* tp_dealloc */
415 0, /* tp_print */
416 0, /* tp_getattr */
417 0, /* tp_setattr */
418 0, /* tp_reserved */
419 (reprfunc)range_repr, /* tp_repr */
420 0, /* tp_as_number */
421 &range_as_sequence, /* tp_as_sequence */
422 0, /* tp_as_mapping */
423 0, /* tp_hash */
424 0, /* tp_call */
425 0, /* tp_str */
426 PyObject_GenericGetAttr, /* tp_getattro */
427 0, /* tp_setattro */
428 0, /* tp_as_buffer */
429 Py_TPFLAGS_DEFAULT, /* tp_flags */
430 range_doc, /* tp_doc */
431 0, /* tp_traverse */
432 0, /* tp_clear */
433 0, /* tp_richcompare */
434 0, /* tp_weaklistoffset */
435 range_iter, /* tp_iter */
436 0, /* tp_iternext */
437 range_methods, /* tp_methods */
438 0, /* tp_members */
439 0, /* tp_getset */
440 0, /* tp_base */
441 0, /* tp_dict */
442 0, /* tp_descr_get */
443 0, /* tp_descr_set */
444 0, /* tp_dictoffset */
445 0, /* tp_init */
446 0, /* tp_alloc */
447 range_new, /* tp_new */
Guido van Rossum12d12c51993-10-26 17:58:25 +0000448};
Raymond Hettinger48165d42002-06-05 20:08:48 +0000449
Walter Dörwald4ad94212007-05-21 18:01:17 +0000450/*********************** range Iterator **************************/
Raymond Hettinger48165d42002-06-05 20:08:48 +0000451
Guido van Rossum805365e2007-05-07 22:24:25 +0000452/* There are 2 types of iterators, one for C longs, the other for
453 Python longs (ie, PyObjects). This should make iteration fast
454 in the normal case, but possible for any numeric value.
455*/
456
Raymond Hettinger48165d42002-06-05 20:08:48 +0000457typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000458 PyObject_HEAD
459 long index;
460 long start;
461 long step;
462 long len;
Raymond Hettinger48165d42002-06-05 20:08:48 +0000463} rangeiterobject;
464
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000465static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000466rangeiter_next(rangeiterobject *r) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000467 if (r->index < r->len)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000468 /* cast to unsigned to avoid possible signed overflow
Mark Dickinsonb43dbc22009-11-15 12:56:08 +0000469 in intermediate calculations. */
470 return PyLong_FromLong((long)(r->start +
471 (unsigned long)(r->index++) * r->step));
Guido van Rossum805365e2007-05-07 22:24:25 +0000472 return NULL;
Raymond Hettinger48165d42002-06-05 20:08:48 +0000473}
474
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000475static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000476rangeiter_len(rangeiterobject *r) {
Christian Heimes217cfd12007-12-02 14:31:20 +0000477 return PyLong_FromLong(r->len - r->index);
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000478}
479
Guido van Rossum805365e2007-05-07 22:24:25 +0000480typedef struct {
481 PyObject_HEAD
482 PyObject *index;
483 PyObject *start;
484 PyObject *step;
485 PyObject *len;
486} longrangeiterobject;
487
488static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000489longrangeiter_len(longrangeiterobject *r, PyObject *no_args) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000490 return PyNumber_Subtract(r->len, r->index);
491}
Guido van Rossum317e7742007-05-08 15:18:31 +0000492
Guido van Rossum805365e2007-05-07 22:24:25 +0000493static PyObject *rangeiter_new(PyTypeObject *, PyObject *args, PyObject *kw);
494
495PyDoc_STRVAR(length_hint_doc,
496 "Private method returning an estimate of len(list(it)).");
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000497
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000498static PyMethodDef rangeiter_methods[] = {
Guido van Rossum805365e2007-05-07 22:24:25 +0000499 {"__length_hint__", (PyCFunction)rangeiter_len, METH_NOARGS,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000500 length_hint_doc},
501 {NULL, NULL} /* sentinel */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000502};
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000503
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000504PyTypeObject PyRangeIter_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000505 PyVarObject_HEAD_INIT(&PyType_Type, 0)
506 "range_iterator", /* tp_name */
507 sizeof(rangeiterobject), /* tp_basicsize */
508 0, /* tp_itemsize */
509 /* methods */
510 (destructor)PyObject_Del, /* tp_dealloc */
511 0, /* tp_print */
512 0, /* tp_getattr */
513 0, /* tp_setattr */
514 0, /* tp_reserved */
515 0, /* tp_repr */
516 0, /* tp_as_number */
517 0, /* tp_as_sequence */
518 0, /* tp_as_mapping */
519 0, /* tp_hash */
520 0, /* tp_call */
521 0, /* tp_str */
522 PyObject_GenericGetAttr, /* tp_getattro */
523 0, /* tp_setattro */
524 0, /* tp_as_buffer */
525 Py_TPFLAGS_DEFAULT, /* tp_flags */
526 0, /* tp_doc */
527 0, /* tp_traverse */
528 0, /* tp_clear */
529 0, /* tp_richcompare */
530 0, /* tp_weaklistoffset */
531 PyObject_SelfIter, /* tp_iter */
532 (iternextfunc)rangeiter_next, /* tp_iternext */
533 rangeiter_methods, /* tp_methods */
534 0, /* tp_members */
535 0, /* tp_getset */
536 0, /* tp_base */
537 0, /* tp_dict */
538 0, /* tp_descr_get */
539 0, /* tp_descr_set */
540 0, /* tp_dictoffset */
541 0, /* tp_init */
542 0, /* tp_alloc */
543 rangeiter_new, /* tp_new */
Guido van Rossum805365e2007-05-07 22:24:25 +0000544};
545
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000546/* Return number of items in range (lo, hi, step). step != 0
547 * required. The result always fits in an unsigned long.
Guido van Rossum805365e2007-05-07 22:24:25 +0000548 */
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000549static unsigned long
Benjamin Peterson155614b2010-11-20 22:44:32 +0000550get_len_of_range(long lo, long hi, long step) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000551 /* -------------------------------------------------------------
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000552 If step > 0 and lo >= hi, or step < 0 and lo <= hi, the range is empty.
553 Else for step > 0, if n values are in the range, the last one is
Guido van Rossum805365e2007-05-07 22:24:25 +0000554 lo + (n-1)*step, which must be <= hi-1. Rearranging,
555 n <= (hi - lo - 1)/step + 1, so taking the floor of the RHS gives
556 the proper value. Since lo < hi in this case, hi-lo-1 >= 0, so
557 the RHS is non-negative and so truncation is the same as the
558 floor. Letting M be the largest positive long, the worst case
559 for the RHS numerator is hi=M, lo=-M-1, and then
560 hi-lo-1 = M-(-M-1)-1 = 2*M. Therefore unsigned long has enough
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000561 precision to compute the RHS exactly. The analysis for step < 0
562 is similar.
Guido van Rossum805365e2007-05-07 22:24:25 +0000563 ---------------------------------------------------------------*/
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000564 assert(step != 0);
565 if (step > 0 && lo < hi)
566 return 1UL + (hi - 1UL - lo) / step;
567 else if (step < 0 && lo > hi)
568 return 1UL + (lo - 1UL - hi) / (0UL - step);
569 else
570 return 0UL;
Guido van Rossum805365e2007-05-07 22:24:25 +0000571}
572
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000573/* Initialize a rangeiter object. If the length of the rangeiter object
574 is not representable as a C long, OverflowError is raised. */
575
Guido van Rossum805365e2007-05-07 22:24:25 +0000576static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000577int_range_iter(long start, long stop, long step) {
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000578 rangeiterobject *it = PyObject_New(rangeiterobject, &PyRangeIter_Type);
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000579 unsigned long ulen;
Guido van Rossum805365e2007-05-07 22:24:25 +0000580 if (it == NULL)
581 return NULL;
582 it->start = start;
583 it->step = step;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000584 ulen = get_len_of_range(start, stop, step);
585 if (ulen > (unsigned long)LONG_MAX) {
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000586 Py_DECREF(it);
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000587 PyErr_SetString(PyExc_OverflowError,
588 "range too large to represent as a range_iterator");
589 return NULL;
590 }
591 it->len = (long)ulen;
Guido van Rossum805365e2007-05-07 22:24:25 +0000592 it->index = 0;
593 return (PyObject *)it;
594}
595
596static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000597rangeiter_new(PyTypeObject *type, PyObject *args, PyObject *kw) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000598 long start, stop, step;
599
600 if (!_PyArg_NoKeywords("rangeiter()", kw))
601 return NULL;
602
603 if (!PyArg_ParseTuple(args, "lll;rangeiter() requires 3 int arguments",
604 &start, &stop, &step))
605 return NULL;
606
607 return int_range_iter(start, stop, step);
608}
609
610static PyMethodDef longrangeiter_methods[] = {
611 {"__length_hint__", (PyCFunction)longrangeiter_len, METH_NOARGS,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000612 length_hint_doc},
613 {NULL, NULL} /* sentinel */
Guido van Rossum805365e2007-05-07 22:24:25 +0000614};
615
616static void
Benjamin Peterson155614b2010-11-20 22:44:32 +0000617longrangeiter_dealloc(longrangeiterobject *r) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000618 Py_XDECREF(r->index);
Guido van Rossum317e7742007-05-08 15:18:31 +0000619 Py_XDECREF(r->start);
620 Py_XDECREF(r->step);
621 Py_XDECREF(r->len);
Amaury Forgeot d'Arcb7f17e42007-11-15 20:52:21 +0000622 PyObject_Del(r);
Guido van Rossum805365e2007-05-07 22:24:25 +0000623}
624
625static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000626longrangeiter_next(longrangeiterobject *r) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000627 PyObject *one, *product, *new_index, *result;
628 if (PyObject_RichCompareBool(r->index, r->len, Py_LT) != 1)
629 return NULL;
630
631 one = PyLong_FromLong(1);
632 if (!one)
633 return NULL;
634
Guido van Rossum805365e2007-05-07 22:24:25 +0000635 new_index = PyNumber_Add(r->index, one);
636 Py_DECREF(one);
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000637 if (!new_index)
638 return NULL;
639
640 product = PyNumber_Multiply(r->index, r->step);
641 if (!product) {
642 Py_DECREF(new_index);
Guido van Rossum805365e2007-05-07 22:24:25 +0000643 return NULL;
644 }
645
646 result = PyNumber_Add(r->start, product);
647 Py_DECREF(product);
648 if (result) {
649 Py_DECREF(r->index);
650 r->index = new_index;
651 }
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000652 else {
653 Py_DECREF(new_index);
654 }
Guido van Rossum805365e2007-05-07 22:24:25 +0000655
656 return result;
657}
658
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000659PyTypeObject PyLongRangeIter_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000660 PyVarObject_HEAD_INIT(&PyType_Type, 0)
661 "longrange_iterator", /* tp_name */
662 sizeof(longrangeiterobject), /* tp_basicsize */
663 0, /* tp_itemsize */
664 /* methods */
665 (destructor)longrangeiter_dealloc, /* tp_dealloc */
666 0, /* tp_print */
667 0, /* tp_getattr */
668 0, /* tp_setattr */
669 0, /* tp_reserved */
670 0, /* tp_repr */
671 0, /* tp_as_number */
672 0, /* tp_as_sequence */
673 0, /* tp_as_mapping */
674 0, /* tp_hash */
675 0, /* tp_call */
676 0, /* tp_str */
677 PyObject_GenericGetAttr, /* tp_getattro */
678 0, /* tp_setattro */
679 0, /* tp_as_buffer */
680 Py_TPFLAGS_DEFAULT, /* tp_flags */
681 0, /* tp_doc */
682 0, /* tp_traverse */
683 0, /* tp_clear */
684 0, /* tp_richcompare */
685 0, /* tp_weaklistoffset */
686 PyObject_SelfIter, /* tp_iter */
687 (iternextfunc)longrangeiter_next, /* tp_iternext */
688 longrangeiter_methods, /* tp_methods */
689 0,
Raymond Hettinger48165d42002-06-05 20:08:48 +0000690};
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000691
692static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000693range_iter(PyObject *seq) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000694 rangeobject *r = (rangeobject *)seq;
695 longrangeiterobject *it;
Martin v. Löwis84451042007-12-20 22:57:23 +0000696 long lstart, lstop, lstep;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000697 PyObject *int_it;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000698
Guido van Rossum805365e2007-05-07 22:24:25 +0000699 assert(PyRange_Check(seq));
Martin v. Löwis84451042007-12-20 22:57:23 +0000700
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000701 /* If all three fields and the length convert to long, use the int
702 * version */
Martin v. Löwis84451042007-12-20 22:57:23 +0000703 lstart = PyLong_AsLong(r->start);
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000704 if (lstart == -1 && PyErr_Occurred()) {
705 PyErr_Clear();
706 goto long_range;
Martin v. Löwis84451042007-12-20 22:57:23 +0000707 }
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000708 lstop = PyLong_AsLong(r->stop);
709 if (lstop == -1 && PyErr_Occurred()) {
710 PyErr_Clear();
711 goto long_range;
712 }
713 lstep = PyLong_AsLong(r->step);
714 if (lstep == -1 && PyErr_Occurred()) {
715 PyErr_Clear();
716 goto long_range;
717 }
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000718 int_it = int_range_iter(lstart, lstop, lstep);
719 if (int_it == NULL && PyErr_ExceptionMatches(PyExc_OverflowError)) {
720 PyErr_Clear();
721 goto long_range;
722 }
723 return (PyObject *)int_it;
724
725 long_range:
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000726 it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
Guido van Rossum805365e2007-05-07 22:24:25 +0000727 if (it == NULL)
728 return NULL;
Guido van Rossum317e7742007-05-08 15:18:31 +0000729
730 /* Do all initialization here, so we can DECREF on failure. */
Guido van Rossum805365e2007-05-07 22:24:25 +0000731 it->start = r->start;
Guido van Rossum317e7742007-05-08 15:18:31 +0000732 it->step = r->step;
733 Py_INCREF(it->start);
734 Py_INCREF(it->step);
735
736 it->len = it->index = NULL;
737
Mark Dickinsoneb36d312009-06-24 18:36:46 +0000738 it->len = range_length_obj(r);
739 if (!it->len)
Guido van Rossum805365e2007-05-07 22:24:25 +0000740 goto create_failure;
Guido van Rossum805365e2007-05-07 22:24:25 +0000741 it->index = PyLong_FromLong(0);
742 if (!it->index)
743 goto create_failure;
744
Guido van Rossum805365e2007-05-07 22:24:25 +0000745 return (PyObject *)it;
746
747create_failure:
Guido van Rossum317e7742007-05-08 15:18:31 +0000748 Py_DECREF(it);
Guido van Rossum805365e2007-05-07 22:24:25 +0000749 return NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000750}
751
752static PyObject *
Benjamin Peterson155614b2010-11-20 22:44:32 +0000753range_reverse(PyObject *seq) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000754 rangeobject *range = (rangeobject*) seq;
755 longrangeiterobject *it;
756 PyObject *one, *sum, *diff, *len = NULL, *product;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000757 long lstart, lstop, lstep, new_start, new_stop;
758 unsigned long ulen;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000759
Guido van Rossum805365e2007-05-07 22:24:25 +0000760 assert(PyRange_Check(seq));
Martin v. Löwis84451042007-12-20 22:57:23 +0000761
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000762 /* reversed(range(start, stop, step)) can be expressed as
763 range(start+(n-1)*step, start-step, -step), where n is the number of
764 integers in the range.
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000765
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000766 If each of start, stop, step, -step, start-step, and the length
767 of the iterator is representable as a C long, use the int
768 version. This excludes some cases where the reversed range is
769 representable as a range_iterator, but it's good enough for
770 common cases and it makes the checks simple. */
771
772 lstart = PyLong_AsLong(range->start);
773 if (lstart == -1 && PyErr_Occurred()) {
774 PyErr_Clear();
775 goto long_range;
776 }
777 lstop = PyLong_AsLong(range->stop);
778 if (lstop == -1 && PyErr_Occurred()) {
779 PyErr_Clear();
780 goto long_range;
781 }
782 lstep = PyLong_AsLong(range->step);
783 if (lstep == -1 && PyErr_Occurred()) {
784 PyErr_Clear();
785 goto long_range;
786 }
787 /* check for possible overflow of -lstep */
788 if (lstep == LONG_MIN)
789 goto long_range;
790
791 /* check for overflow of lstart - lstep:
792
793 for lstep > 0, need only check whether lstart - lstep < LONG_MIN.
794 for lstep < 0, need only check whether lstart - lstep > LONG_MAX
795
796 Rearrange these inequalities as:
797
798 lstart - LONG_MIN < lstep (lstep > 0)
799 LONG_MAX - lstart < -lstep (lstep < 0)
800
801 and compute both sides as unsigned longs, to avoid the
802 possibility of undefined behaviour due to signed overflow. */
803
804 if (lstep > 0) {
805 if ((unsigned long)lstart - LONG_MIN < (unsigned long)lstep)
806 goto long_range;
807 }
808 else {
809 if (LONG_MAX - (unsigned long)lstart < 0UL - lstep)
810 goto long_range;
811 }
812
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000813 ulen = get_len_of_range(lstart, lstop, lstep);
814 if (ulen > (unsigned long)LONG_MAX)
815 goto long_range;
816
817 new_stop = lstart - lstep;
818 new_start = (long)(new_stop + ulen * lstep);
819 return int_range_iter(new_start, new_stop, -lstep);
820
821long_range:
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000822 it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
Guido van Rossum805365e2007-05-07 22:24:25 +0000823 if (it == NULL)
824 return NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000825
Guido van Rossum805365e2007-05-07 22:24:25 +0000826 /* start + (len - 1) * step */
827 len = range_length_obj(range);
828 if (!len)
829 goto create_failure;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000830
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000831 /* Steal reference to len. */
832 it->len = len;
833
Guido van Rossum805365e2007-05-07 22:24:25 +0000834 one = PyLong_FromLong(1);
835 if (!one)
836 goto create_failure;
837
838 diff = PyNumber_Subtract(len, one);
839 Py_DECREF(one);
840 if (!diff)
841 goto create_failure;
842
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000843 product = PyNumber_Multiply(diff, range->step);
844 Py_DECREF(diff);
Guido van Rossum805365e2007-05-07 22:24:25 +0000845 if (!product)
846 goto create_failure;
847
848 sum = PyNumber_Add(range->start, product);
849 Py_DECREF(product);
850 it->start = sum;
851 if (!it->start)
852 goto create_failure;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000853
Guido van Rossum805365e2007-05-07 22:24:25 +0000854 it->step = PyNumber_Negative(range->step);
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000855 if (!it->step)
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000856 goto create_failure;
Guido van Rossum805365e2007-05-07 22:24:25 +0000857
858 it->index = PyLong_FromLong(0);
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000859 if (!it->index)
860 goto create_failure;
Guido van Rossum805365e2007-05-07 22:24:25 +0000861
862 return (PyObject *)it;
863
864create_failure:
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000865 Py_DECREF(it);
Guido van Rossum805365e2007-05-07 22:24:25 +0000866 return NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000867}