blob: 6ce0187adbeca26b4296095c827e861f9a2b90b0 [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 Petersona1864f32010-11-20 23:05:39 +000023validate_step(PyObject *step)
24{
Guido van Rossum805365e2007-05-07 22:24:25 +000025 /* No step specified, use a step of 1. */
26 if (!step)
Christian Heimes217cfd12007-12-02 14:31:20 +000027 return PyLong_FromLong(1);
Guido van Rossum805365e2007-05-07 22:24:25 +000028
29 step = PyNumber_Index(step);
30 if (step) {
31 Py_ssize_t istep = PyNumber_AsSsize_t(step, NULL);
32 if (istep == -1 && PyErr_Occurred()) {
33 /* Ignore OverflowError, we know the value isn't 0. */
34 PyErr_Clear();
35 }
36 else if (istep == 0) {
37 PyErr_SetString(PyExc_ValueError,
38 "range() arg 3 must not be zero");
39 Py_CLEAR(step);
40 }
41 }
42
43 return step;
Raymond Hettingerc4c453f2002-06-05 23:12:45 +000044}
45
Guido van Rossum805365e2007-05-07 22:24:25 +000046/* XXX(nnorwitz): should we error check if the user passes any empty ranges?
47 range(-10)
48 range(0, -5)
49 range(0, 5, -1)
50*/
Raymond Hettingerc4c453f2002-06-05 23:12:45 +000051static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +000052range_new(PyTypeObject *type, PyObject *args, PyObject *kw)
53{
Guido van Rossum805365e2007-05-07 22:24:25 +000054 rangeobject *obj = NULL;
55 PyObject *start = NULL, *stop = NULL, *step = NULL;
Raymond Hettingerc4c453f2002-06-05 23:12:45 +000056
Guido van Rossum805365e2007-05-07 22:24:25 +000057 if (!_PyArg_NoKeywords("range()", kw))
58 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +000059
Guido van Rossum805365e2007-05-07 22:24:25 +000060 if (PyTuple_Size(args) <= 1) {
61 if (!PyArg_UnpackTuple(args, "range", 1, 1, &stop))
Raymond Hettinger94f55832009-06-12 18:40:16 +000062 return NULL;
Guido van Rossum805365e2007-05-07 22:24:25 +000063 stop = PyNumber_Index(stop);
64 if (!stop)
Raymond Hettinger94f55832009-06-12 18:40:16 +000065 return NULL;
Christian Heimes217cfd12007-12-02 14:31:20 +000066 start = PyLong_FromLong(0);
Raymond Hettinger94f55832009-06-12 18:40:16 +000067 if (!start) {
68 Py_DECREF(stop);
69 return NULL;
70 }
Christian Heimes217cfd12007-12-02 14:31:20 +000071 step = PyLong_FromLong(1);
Raymond Hettinger94f55832009-06-12 18:40:16 +000072 if (!step) {
73 Py_DECREF(stop);
74 Py_DECREF(start);
75 return NULL;
76 }
Guido van Rossum805365e2007-05-07 22:24:25 +000077 }
78 else {
79 if (!PyArg_UnpackTuple(args, "range", 2, 3,
80 &start, &stop, &step))
Raymond Hettinger94f55832009-06-12 18:40:16 +000081 return NULL;
Tim Petersfeec4532004-08-08 07:17:39 +000082
Guido van Rossum805365e2007-05-07 22:24:25 +000083 /* Convert borrowed refs to owned refs */
84 start = PyNumber_Index(start);
Raymond Hettinger94f55832009-06-12 18:40:16 +000085 if (!start)
86 return NULL;
Guido van Rossum805365e2007-05-07 22:24:25 +000087 stop = PyNumber_Index(stop);
Raymond Hettinger94f55832009-06-12 18:40:16 +000088 if (!stop) {
89 Py_DECREF(start);
90 return NULL;
91 }
92 step = validate_step(step); /* Caution, this can clear exceptions */
93 if (!step) {
94 Py_DECREF(start);
95 Py_DECREF(stop);
96 return NULL;
97 }
Guido van Rossum805365e2007-05-07 22:24:25 +000098 }
99
100 obj = PyObject_New(rangeobject, &PyRange_Type);
101 if (obj == NULL)
102 goto Fail;
103 obj->start = start;
104 obj->stop = stop;
105 obj->step = step;
106 return (PyObject *) obj;
107
108Fail:
109 Py_XDECREF(start);
110 Py_XDECREF(stop);
111 Py_XDECREF(step);
112 return NULL;
Raymond Hettingerc4c453f2002-06-05 23:12:45 +0000113}
114
Martin v. Löwis14f8b4c2002-06-13 20:33:02 +0000115PyDoc_STRVAR(range_doc,
Guido van Rossum805365e2007-05-07 22:24:25 +0000116"range([start,] stop[, step]) -> range object\n\
Raymond Hettingerc4c453f2002-06-05 23:12:45 +0000117\n\
Guido van Rossum805365e2007-05-07 22:24:25 +0000118Returns an iterator that generates the numbers in the range on demand.");
Raymond Hettingerc4c453f2002-06-05 23:12:45 +0000119
Guido van Rossum805365e2007-05-07 22:24:25 +0000120static void
Benjamin Petersona1864f32010-11-20 23:05:39 +0000121range_dealloc(rangeobject *r)
122{
Guido van Rossum805365e2007-05-07 22:24:25 +0000123 Py_DECREF(r->start);
124 Py_DECREF(r->stop);
125 Py_DECREF(r->step);
Amaury Forgeot d'Arcb7f17e42007-11-15 20:52:21 +0000126 PyObject_Del(r);
Guido van Rossum805365e2007-05-07 22:24:25 +0000127}
128
Mark Dickinson732166d2009-06-28 22:08:40 +0000129/* Return number of items in range (lo, hi, step) as a PyLong object,
130 * when arguments are PyLong objects. Arguments MUST return 1 with
131 * PyLong_Check(). Return NULL when there is an error.
Guido van Rossum805365e2007-05-07 22:24:25 +0000132 */
133static PyObject*
Benjamin Petersona1864f32010-11-20 23:05:39 +0000134range_length_obj(rangeobject *r)
135{
Guido van Rossum805365e2007-05-07 22:24:25 +0000136 /* -------------------------------------------------------------
137 Algorithm is equal to that of get_len_of_range(), but it operates
Benjamin Petersona47af9c2009-06-24 21:14:38 +0000138 on PyObjects (which are assumed to be PyLong objects).
Guido van Rossum805365e2007-05-07 22:24:25 +0000139 ---------------------------------------------------------------*/
Mark Dickinson211c6252009-02-01 10:28:51 +0000140 int cmp_result;
Guido van Rossum805365e2007-05-07 22:24:25 +0000141 PyObject *lo, *hi;
142 PyObject *step = NULL;
143 PyObject *diff = NULL;
144 PyObject *one = NULL;
145 PyObject *tmp1 = NULL, *tmp2 = NULL, *result;
146 /* holds sub-expression evaluations */
147
148 PyObject *zero = PyLong_FromLong(0);
149 if (zero == NULL)
150 return NULL;
Mark Dickinson211c6252009-02-01 10:28:51 +0000151 cmp_result = PyObject_RichCompareBool(r->step, zero, Py_GT);
Guido van Rossum805365e2007-05-07 22:24:25 +0000152 Py_DECREF(zero);
Mark Dickinson211c6252009-02-01 10:28:51 +0000153 if (cmp_result == -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000154 return NULL;
155
Mark Dickinson211c6252009-02-01 10:28:51 +0000156 if (cmp_result == 1) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000157 lo = r->start;
158 hi = r->stop;
159 step = r->step;
160 Py_INCREF(step);
161 } else {
162 lo = r->stop;
163 hi = r->start;
164 step = PyNumber_Negative(r->step);
165 if (!step)
166 return NULL;
167 }
168
169 /* if (lo >= hi), return length of 0. */
Mark Dickinson211c6252009-02-01 10:28:51 +0000170 if (PyObject_RichCompareBool(lo, hi, Py_GE) == 1) {
Guido van Rossum805365e2007-05-07 22:24:25 +0000171 Py_XDECREF(step);
172 return PyLong_FromLong(0);
173 }
174
175 if ((one = PyLong_FromLong(1L)) == NULL)
176 goto Fail;
177
178 if ((tmp1 = PyNumber_Subtract(hi, lo)) == NULL)
179 goto Fail;
180
181 if ((diff = PyNumber_Subtract(tmp1, one)) == NULL)
182 goto Fail;
183
184 if ((tmp2 = PyNumber_FloorDivide(diff, step)) == NULL)
185 goto Fail;
186
187 if ((result = PyNumber_Add(tmp2, one)) == NULL)
188 goto Fail;
189
190 Py_DECREF(tmp2);
191 Py_DECREF(diff);
192 Py_DECREF(step);
193 Py_DECREF(tmp1);
194 Py_DECREF(one);
195 return result;
196
197 Fail:
198 Py_XDECREF(tmp2);
199 Py_XDECREF(diff);
200 Py_XDECREF(step);
201 Py_XDECREF(tmp1);
202 Py_XDECREF(one);
203 return NULL;
Guido van Rossum12d12c51993-10-26 17:58:25 +0000204}
205
Martin v. Löwis18e16552006-02-15 17:27:45 +0000206static Py_ssize_t
Benjamin Petersona1864f32010-11-20 23:05:39 +0000207range_length(rangeobject *r)
208{
Guido van Rossum805365e2007-05-07 22:24:25 +0000209 PyObject *len = range_length_obj(r);
210 Py_ssize_t result = -1;
211 if (len) {
212 result = PyLong_AsSsize_t(len);
213 Py_DECREF(len);
214 }
215 return result;
216}
217
218/* range(...)[x] is necessary for: seq[:] = range(...) */
219
220static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000221range_item(rangeobject *r, Py_ssize_t i)
222{
Guido van Rossum805365e2007-05-07 22:24:25 +0000223 Py_ssize_t len = range_length(r);
224 PyObject *rem, *incr, *result;
225
226 /* XXX(nnorwitz): should negative indices be supported? */
227 /* XXX(nnorwitz): should support range[x] where x > PY_SSIZE_T_MAX? */
228 if (i < 0 || i >= len) {
229 if (!PyErr_Occurred())
230 PyErr_SetString(PyExc_IndexError,
Walter Dörwald4ad94212007-05-21 18:01:17 +0000231 "range object index out of range");
Benjamin Petersondf0a5cb2008-04-25 21:15:37 +0000232 return NULL;
233 }
Guido van Rossum805365e2007-05-07 22:24:25 +0000234
235 /* XXX(nnorwitz): optimize for short ints. */
Raymond Hettingerad3f3322008-02-09 04:13:49 +0000236 rem = PyLong_FromSsize_t(i);
Guido van Rossum805365e2007-05-07 22:24:25 +0000237 if (!rem)
238 return NULL;
239 incr = PyNumber_Multiply(rem, r->step);
240 Py_DECREF(rem);
241 if (!incr)
242 return NULL;
243 result = PyNumber_Add(r->start, incr);
244 Py_DECREF(incr);
245 return result;
Guido van Rossum7d6aa511993-12-21 22:50:31 +0000246}
247
Guido van Rossumc0b618a1997-05-02 03:12:38 +0000248static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000249range_repr(rangeobject *r)
250{
Walter Dörwald03b43d82007-05-21 10:43:34 +0000251 Py_ssize_t istep;
Tim Petersd976ab72004-08-08 06:29:10 +0000252
Guido van Rossum805365e2007-05-07 22:24:25 +0000253 /* Check for special case values for printing. We don't always
Walter Dörwald03b43d82007-05-21 10:43:34 +0000254 need the step value. We don't care about errors
Guido van Rossum805365e2007-05-07 22:24:25 +0000255 (it means overflow), so clear the errors. */
Guido van Rossum805365e2007-05-07 22:24:25 +0000256 istep = PyNumber_AsSsize_t(r->step, NULL);
257 if (istep != 1 || (istep == -1 && PyErr_Occurred())) {
258 PyErr_Clear();
Guido van Rossum805365e2007-05-07 22:24:25 +0000259 }
260
Walter Dörwald03b43d82007-05-21 10:43:34 +0000261 if (istep == 1)
Walter Dörwald850e5162007-05-20 08:19:54 +0000262 return PyUnicode_FromFormat("range(%R, %R)", r->start, r->stop);
263 else
264 return PyUnicode_FromFormat("range(%R, %R, %R)",
265 r->start, r->stop, r->step);
Thomas Woutersefafcea2001-07-09 12:30:54 +0000266}
267
Alexandre Vassalotti75056072008-06-10 04:03:04 +0000268/* Pickling support */
269static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000270range_reduce(rangeobject *r, PyObject *args)
271{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000272 return Py_BuildValue("(O(OOO))", Py_TYPE(r),
Alexandre Vassalotti75056072008-06-10 04:03:04 +0000273 r->start, r->stop, r->step);
274}
275
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000276/* Assumes (PyLong_CheckExact(ob) || PyBool_Check(ob)) */
277static int
Benjamin Petersona1864f32010-11-20 23:05:39 +0000278range_contains_long(rangeobject *r, PyObject *ob)
279{
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000280 int cmp1, cmp2, cmp3;
281 PyObject *tmp1 = NULL;
282 PyObject *tmp2 = NULL;
283 PyObject *zero = NULL;
284 int result = -1;
285
286 zero = PyLong_FromLong(0);
287 if (zero == NULL) /* MemoryError in int(0) */
288 goto end;
289
290 /* Check if the value can possibly be in the range. */
291
292 cmp1 = PyObject_RichCompareBool(r->step, zero, Py_GT);
293 if (cmp1 == -1)
294 goto end;
295 if (cmp1 == 1) { /* positive steps: start <= ob < stop */
296 cmp2 = PyObject_RichCompareBool(r->start, ob, Py_LE);
297 cmp3 = PyObject_RichCompareBool(ob, r->stop, Py_LT);
298 }
299 else { /* negative steps: stop < ob <= start */
300 cmp2 = PyObject_RichCompareBool(ob, r->start, Py_LE);
301 cmp3 = PyObject_RichCompareBool(r->stop, ob, Py_LT);
302 }
303
304 if (cmp2 == -1 || cmp3 == -1) /* TypeError */
305 goto end;
306 if (cmp2 == 0 || cmp3 == 0) { /* ob outside of range */
307 result = 0;
308 goto end;
309 }
310
311 /* Check that the stride does not invalidate ob's membership. */
312 tmp1 = PyNumber_Subtract(ob, r->start);
313 if (tmp1 == NULL)
314 goto end;
315 tmp2 = PyNumber_Remainder(tmp1, r->step);
316 if (tmp2 == NULL)
317 goto end;
318 /* result = (int(ob) - start % step) == 0 */
319 result = PyObject_RichCompareBool(tmp2, zero, Py_EQ);
320 end:
321 Py_XDECREF(tmp1);
322 Py_XDECREF(tmp2);
323 Py_XDECREF(zero);
324 return result;
325}
326
Mark Dickinson3e124ae2009-09-22 21:47:24 +0000327static int
Benjamin Petersona1864f32010-11-20 23:05:39 +0000328range_contains(rangeobject *r, PyObject *ob)
329{
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000330 if (PyLong_CheckExact(ob) || PyBool_Check(ob))
331 return range_contains_long(r, ob);
Mark Dickinson3e124ae2009-09-22 21:47:24 +0000332
Mark Dickinson3e124ae2009-09-22 21:47:24 +0000333 return (int)_PySequence_IterSearch((PyObject*)r, ob,
334 PY_ITERSEARCH_CONTAINS);
335}
336
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000337static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000338range_count(rangeobject *r, PyObject *ob)
339{
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000340 if (PyLong_CheckExact(ob) || PyBool_Check(ob)) {
Georg Brandl7e5343b2010-11-20 22:40:10 +0000341 int result = range_contains_long(r, ob);
342 if (result == -1)
343 return NULL;
344 else if (result)
Benjamin Peterson0b458d52010-11-20 22:35:41 +0000345 return PyLong_FromLong(1);
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000346 else
Benjamin Peterson0b458d52010-11-20 22:35:41 +0000347 return PyLong_FromLong(0);
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000348 } else {
349 Py_ssize_t count;
350 count = _PySequence_IterSearch((PyObject*)r, ob, PY_ITERSEARCH_COUNT);
351 if (count == -1)
352 return NULL;
353 return PyLong_FromSsize_t(count);
354 }
355}
356
357static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000358range_index(rangeobject *r, PyObject *ob)
359{
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000360 int contains;
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000361
362 if (!PyLong_CheckExact(ob) && !PyBool_Check(ob)) {
363 Py_ssize_t index;
364 index = _PySequence_IterSearch((PyObject*)r, ob, PY_ITERSEARCH_INDEX);
365 if (index == -1)
366 return NULL;
367 return PyLong_FromSsize_t(index);
368 }
369
370 contains = range_contains_long(r, ob);
371 if (contains == -1)
372 return NULL;
373
Benjamin Peterson155614b2010-11-20 22:44:32 +0000374 if (contains) {
375 PyObject *idx, *tmp = PyNumber_Subtract(ob, r->start);
376 if (tmp == NULL)
377 return NULL;
378 /* idx = (ob - r.start) // r.step */
379 idx = PyNumber_FloorDivide(tmp, r->step);
380 Py_DECREF(tmp);
381 return idx;
382 }
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000383
384 /* object is not in the range */
Benjamin Peterson155614b2010-11-20 22:44:32 +0000385 PyErr_Format(PyExc_ValueError, "%R is not in range", ob);
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000386 return NULL;
387}
388
Guido van Rossumc0b618a1997-05-02 03:12:38 +0000389static PySequenceMethods range_as_sequence = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000390 (lenfunc)range_length, /* sq_length */
391 0, /* sq_concat */
392 0, /* sq_repeat */
Guido van Rossum805365e2007-05-07 22:24:25 +0000393 (ssizeargfunc)range_item, /* sq_item */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000394 0, /* sq_slice */
Mark Dickinson3e124ae2009-09-22 21:47:24 +0000395 0, /* sq_ass_item */
396 0, /* sq_ass_slice */
397 (objobjproc)range_contains, /* sq_contains */
Guido van Rossum12d12c51993-10-26 17:58:25 +0000398};
399
Jeremy Hylton938ace62002-07-17 16:30:39 +0000400static PyObject * range_iter(PyObject *seq);
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000401static PyObject * range_reverse(PyObject *seq);
402
403PyDoc_STRVAR(reverse_doc,
404"Returns a reverse iterator.");
405
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000406PyDoc_STRVAR(count_doc,
407"rangeobject.count(value) -> integer -- return number of occurrences of value");
408
409PyDoc_STRVAR(index_doc,
410"rangeobject.index(value, [start, [stop]]) -> integer -- return index of value.\n"
411"Raises ValueError if the value is not present.");
412
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000413static PyMethodDef range_methods[] = {
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000414 {"__reversed__", (PyCFunction)range_reverse, METH_NOARGS, reverse_doc},
415 {"__reduce__", (PyCFunction)range_reduce, METH_VARARGS},
416 {"count", (PyCFunction)range_count, METH_O, count_doc},
417 {"index", (PyCFunction)range_index, METH_O, index_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000418 {NULL, NULL} /* sentinel */
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000419};
Raymond Hettinger48165d42002-06-05 20:08:48 +0000420
Guido van Rossumc0b618a1997-05-02 03:12:38 +0000421PyTypeObject PyRange_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000422 PyVarObject_HEAD_INIT(&PyType_Type, 0)
423 "range", /* Name of this type */
424 sizeof(rangeobject), /* Basic object size */
425 0, /* Item size for varobject */
426 (destructor)range_dealloc, /* tp_dealloc */
427 0, /* tp_print */
428 0, /* tp_getattr */
429 0, /* tp_setattr */
430 0, /* tp_reserved */
431 (reprfunc)range_repr, /* tp_repr */
432 0, /* tp_as_number */
433 &range_as_sequence, /* tp_as_sequence */
434 0, /* tp_as_mapping */
435 0, /* tp_hash */
436 0, /* tp_call */
437 0, /* tp_str */
438 PyObject_GenericGetAttr, /* tp_getattro */
439 0, /* tp_setattro */
440 0, /* tp_as_buffer */
441 Py_TPFLAGS_DEFAULT, /* tp_flags */
442 range_doc, /* tp_doc */
443 0, /* tp_traverse */
444 0, /* tp_clear */
445 0, /* tp_richcompare */
446 0, /* tp_weaklistoffset */
447 range_iter, /* tp_iter */
448 0, /* tp_iternext */
449 range_methods, /* tp_methods */
450 0, /* tp_members */
451 0, /* tp_getset */
452 0, /* tp_base */
453 0, /* tp_dict */
454 0, /* tp_descr_get */
455 0, /* tp_descr_set */
456 0, /* tp_dictoffset */
457 0, /* tp_init */
458 0, /* tp_alloc */
459 range_new, /* tp_new */
Guido van Rossum12d12c51993-10-26 17:58:25 +0000460};
Raymond Hettinger48165d42002-06-05 20:08:48 +0000461
Walter Dörwald4ad94212007-05-21 18:01:17 +0000462/*********************** range Iterator **************************/
Raymond Hettinger48165d42002-06-05 20:08:48 +0000463
Guido van Rossum805365e2007-05-07 22:24:25 +0000464/* There are 2 types of iterators, one for C longs, the other for
465 Python longs (ie, PyObjects). This should make iteration fast
466 in the normal case, but possible for any numeric value.
467*/
468
Raymond Hettinger48165d42002-06-05 20:08:48 +0000469typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000470 PyObject_HEAD
471 long index;
472 long start;
473 long step;
474 long len;
Raymond Hettinger48165d42002-06-05 20:08:48 +0000475} rangeiterobject;
476
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000477static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000478rangeiter_next(rangeiterobject *r)
479{
Guido van Rossum805365e2007-05-07 22:24:25 +0000480 if (r->index < r->len)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000481 /* cast to unsigned to avoid possible signed overflow
Mark Dickinsonb43dbc22009-11-15 12:56:08 +0000482 in intermediate calculations. */
483 return PyLong_FromLong((long)(r->start +
484 (unsigned long)(r->index++) * r->step));
Guido van Rossum805365e2007-05-07 22:24:25 +0000485 return NULL;
Raymond Hettinger48165d42002-06-05 20:08:48 +0000486}
487
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000488static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000489rangeiter_len(rangeiterobject *r)
490{
Christian Heimes217cfd12007-12-02 14:31:20 +0000491 return PyLong_FromLong(r->len - r->index);
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000492}
493
Guido van Rossum805365e2007-05-07 22:24:25 +0000494typedef struct {
495 PyObject_HEAD
496 PyObject *index;
497 PyObject *start;
498 PyObject *step;
499 PyObject *len;
500} longrangeiterobject;
501
502static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000503longrangeiter_len(longrangeiterobject *r, PyObject *no_args)
504{
Guido van Rossum805365e2007-05-07 22:24:25 +0000505 return PyNumber_Subtract(r->len, r->index);
506}
Guido van Rossum317e7742007-05-08 15:18:31 +0000507
Guido van Rossum805365e2007-05-07 22:24:25 +0000508static PyObject *rangeiter_new(PyTypeObject *, PyObject *args, PyObject *kw);
509
510PyDoc_STRVAR(length_hint_doc,
511 "Private method returning an estimate of len(list(it)).");
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000512
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000513static PyMethodDef rangeiter_methods[] = {
Guido van Rossum805365e2007-05-07 22:24:25 +0000514 {"__length_hint__", (PyCFunction)rangeiter_len, METH_NOARGS,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000515 length_hint_doc},
516 {NULL, NULL} /* sentinel */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000517};
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000518
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000519PyTypeObject PyRangeIter_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000520 PyVarObject_HEAD_INIT(&PyType_Type, 0)
521 "range_iterator", /* tp_name */
522 sizeof(rangeiterobject), /* tp_basicsize */
523 0, /* tp_itemsize */
524 /* methods */
525 (destructor)PyObject_Del, /* tp_dealloc */
526 0, /* tp_print */
527 0, /* tp_getattr */
528 0, /* tp_setattr */
529 0, /* tp_reserved */
530 0, /* tp_repr */
531 0, /* tp_as_number */
532 0, /* tp_as_sequence */
533 0, /* tp_as_mapping */
534 0, /* tp_hash */
535 0, /* tp_call */
536 0, /* tp_str */
537 PyObject_GenericGetAttr, /* tp_getattro */
538 0, /* tp_setattro */
539 0, /* tp_as_buffer */
540 Py_TPFLAGS_DEFAULT, /* tp_flags */
541 0, /* tp_doc */
542 0, /* tp_traverse */
543 0, /* tp_clear */
544 0, /* tp_richcompare */
545 0, /* tp_weaklistoffset */
546 PyObject_SelfIter, /* tp_iter */
547 (iternextfunc)rangeiter_next, /* tp_iternext */
548 rangeiter_methods, /* tp_methods */
549 0, /* tp_members */
550 0, /* tp_getset */
551 0, /* tp_base */
552 0, /* tp_dict */
553 0, /* tp_descr_get */
554 0, /* tp_descr_set */
555 0, /* tp_dictoffset */
556 0, /* tp_init */
557 0, /* tp_alloc */
558 rangeiter_new, /* tp_new */
Guido van Rossum805365e2007-05-07 22:24:25 +0000559};
560
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000561/* Return number of items in range (lo, hi, step). step != 0
562 * required. The result always fits in an unsigned long.
Guido van Rossum805365e2007-05-07 22:24:25 +0000563 */
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000564static unsigned long
Benjamin Petersona1864f32010-11-20 23:05:39 +0000565get_len_of_range(long lo, long hi, long step)
566{
Guido van Rossum805365e2007-05-07 22:24:25 +0000567 /* -------------------------------------------------------------
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000568 If step > 0 and lo >= hi, or step < 0 and lo <= hi, the range is empty.
569 Else for step > 0, if n values are in the range, the last one is
Guido van Rossum805365e2007-05-07 22:24:25 +0000570 lo + (n-1)*step, which must be <= hi-1. Rearranging,
571 n <= (hi - lo - 1)/step + 1, so taking the floor of the RHS gives
572 the proper value. Since lo < hi in this case, hi-lo-1 >= 0, so
573 the RHS is non-negative and so truncation is the same as the
574 floor. Letting M be the largest positive long, the worst case
575 for the RHS numerator is hi=M, lo=-M-1, and then
576 hi-lo-1 = M-(-M-1)-1 = 2*M. Therefore unsigned long has enough
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000577 precision to compute the RHS exactly. The analysis for step < 0
578 is similar.
Guido van Rossum805365e2007-05-07 22:24:25 +0000579 ---------------------------------------------------------------*/
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000580 assert(step != 0);
581 if (step > 0 && lo < hi)
582 return 1UL + (hi - 1UL - lo) / step;
583 else if (step < 0 && lo > hi)
584 return 1UL + (lo - 1UL - hi) / (0UL - step);
585 else
586 return 0UL;
Guido van Rossum805365e2007-05-07 22:24:25 +0000587}
588
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000589/* Initialize a rangeiter object. If the length of the rangeiter object
590 is not representable as a C long, OverflowError is raised. */
591
Guido van Rossum805365e2007-05-07 22:24:25 +0000592static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000593int_range_iter(long start, long stop, long step)
594{
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000595 rangeiterobject *it = PyObject_New(rangeiterobject, &PyRangeIter_Type);
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000596 unsigned long ulen;
Guido van Rossum805365e2007-05-07 22:24:25 +0000597 if (it == NULL)
598 return NULL;
599 it->start = start;
600 it->step = step;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000601 ulen = get_len_of_range(start, stop, step);
602 if (ulen > (unsigned long)LONG_MAX) {
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000603 Py_DECREF(it);
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000604 PyErr_SetString(PyExc_OverflowError,
605 "range too large to represent as a range_iterator");
606 return NULL;
607 }
608 it->len = (long)ulen;
Guido van Rossum805365e2007-05-07 22:24:25 +0000609 it->index = 0;
610 return (PyObject *)it;
611}
612
613static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000614rangeiter_new(PyTypeObject *type, PyObject *args, PyObject *kw)
615{
Guido van Rossum805365e2007-05-07 22:24:25 +0000616 long start, stop, step;
617
618 if (!_PyArg_NoKeywords("rangeiter()", kw))
619 return NULL;
620
621 if (!PyArg_ParseTuple(args, "lll;rangeiter() requires 3 int arguments",
622 &start, &stop, &step))
623 return NULL;
624
625 return int_range_iter(start, stop, step);
626}
627
628static PyMethodDef longrangeiter_methods[] = {
629 {"__length_hint__", (PyCFunction)longrangeiter_len, METH_NOARGS,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000630 length_hint_doc},
631 {NULL, NULL} /* sentinel */
Guido van Rossum805365e2007-05-07 22:24:25 +0000632};
633
634static void
Benjamin Petersona1864f32010-11-20 23:05:39 +0000635longrangeiter_dealloc(longrangeiterobject *r)
636{
Guido van Rossum805365e2007-05-07 22:24:25 +0000637 Py_XDECREF(r->index);
Guido van Rossum317e7742007-05-08 15:18:31 +0000638 Py_XDECREF(r->start);
639 Py_XDECREF(r->step);
640 Py_XDECREF(r->len);
Amaury Forgeot d'Arcb7f17e42007-11-15 20:52:21 +0000641 PyObject_Del(r);
Guido van Rossum805365e2007-05-07 22:24:25 +0000642}
643
644static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000645longrangeiter_next(longrangeiterobject *r)
646{
Guido van Rossum805365e2007-05-07 22:24:25 +0000647 PyObject *one, *product, *new_index, *result;
648 if (PyObject_RichCompareBool(r->index, r->len, Py_LT) != 1)
649 return NULL;
650
651 one = PyLong_FromLong(1);
652 if (!one)
653 return NULL;
654
Guido van Rossum805365e2007-05-07 22:24:25 +0000655 new_index = PyNumber_Add(r->index, one);
656 Py_DECREF(one);
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000657 if (!new_index)
658 return NULL;
659
660 product = PyNumber_Multiply(r->index, r->step);
661 if (!product) {
662 Py_DECREF(new_index);
Guido van Rossum805365e2007-05-07 22:24:25 +0000663 return NULL;
664 }
665
666 result = PyNumber_Add(r->start, product);
667 Py_DECREF(product);
668 if (result) {
669 Py_DECREF(r->index);
670 r->index = new_index;
671 }
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000672 else {
673 Py_DECREF(new_index);
674 }
Guido van Rossum805365e2007-05-07 22:24:25 +0000675
676 return result;
677}
678
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000679PyTypeObject PyLongRangeIter_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000680 PyVarObject_HEAD_INIT(&PyType_Type, 0)
681 "longrange_iterator", /* tp_name */
682 sizeof(longrangeiterobject), /* tp_basicsize */
683 0, /* tp_itemsize */
684 /* methods */
685 (destructor)longrangeiter_dealloc, /* tp_dealloc */
686 0, /* tp_print */
687 0, /* tp_getattr */
688 0, /* tp_setattr */
689 0, /* tp_reserved */
690 0, /* tp_repr */
691 0, /* tp_as_number */
692 0, /* tp_as_sequence */
693 0, /* tp_as_mapping */
694 0, /* tp_hash */
695 0, /* tp_call */
696 0, /* tp_str */
697 PyObject_GenericGetAttr, /* tp_getattro */
698 0, /* tp_setattro */
699 0, /* tp_as_buffer */
700 Py_TPFLAGS_DEFAULT, /* tp_flags */
701 0, /* tp_doc */
702 0, /* tp_traverse */
703 0, /* tp_clear */
704 0, /* tp_richcompare */
705 0, /* tp_weaklistoffset */
706 PyObject_SelfIter, /* tp_iter */
707 (iternextfunc)longrangeiter_next, /* tp_iternext */
708 longrangeiter_methods, /* tp_methods */
709 0,
Raymond Hettinger48165d42002-06-05 20:08:48 +0000710};
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000711
712static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000713range_iter(PyObject *seq)
714{
Guido van Rossum805365e2007-05-07 22:24:25 +0000715 rangeobject *r = (rangeobject *)seq;
716 longrangeiterobject *it;
Martin v. Löwis84451042007-12-20 22:57:23 +0000717 long lstart, lstop, lstep;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000718 PyObject *int_it;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000719
Guido van Rossum805365e2007-05-07 22:24:25 +0000720 assert(PyRange_Check(seq));
Martin v. Löwis84451042007-12-20 22:57:23 +0000721
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000722 /* If all three fields and the length convert to long, use the int
723 * version */
Martin v. Löwis84451042007-12-20 22:57:23 +0000724 lstart = PyLong_AsLong(r->start);
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000725 if (lstart == -1 && PyErr_Occurred()) {
726 PyErr_Clear();
727 goto long_range;
Martin v. Löwis84451042007-12-20 22:57:23 +0000728 }
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000729 lstop = PyLong_AsLong(r->stop);
730 if (lstop == -1 && PyErr_Occurred()) {
731 PyErr_Clear();
732 goto long_range;
733 }
734 lstep = PyLong_AsLong(r->step);
735 if (lstep == -1 && PyErr_Occurred()) {
736 PyErr_Clear();
737 goto long_range;
738 }
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000739 int_it = int_range_iter(lstart, lstop, lstep);
740 if (int_it == NULL && PyErr_ExceptionMatches(PyExc_OverflowError)) {
741 PyErr_Clear();
742 goto long_range;
743 }
744 return (PyObject *)int_it;
745
746 long_range:
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000747 it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
Guido van Rossum805365e2007-05-07 22:24:25 +0000748 if (it == NULL)
749 return NULL;
Guido van Rossum317e7742007-05-08 15:18:31 +0000750
751 /* Do all initialization here, so we can DECREF on failure. */
Guido van Rossum805365e2007-05-07 22:24:25 +0000752 it->start = r->start;
Guido van Rossum317e7742007-05-08 15:18:31 +0000753 it->step = r->step;
754 Py_INCREF(it->start);
755 Py_INCREF(it->step);
756
757 it->len = it->index = NULL;
758
Mark Dickinsoneb36d312009-06-24 18:36:46 +0000759 it->len = range_length_obj(r);
760 if (!it->len)
Guido van Rossum805365e2007-05-07 22:24:25 +0000761 goto create_failure;
Guido van Rossum805365e2007-05-07 22:24:25 +0000762 it->index = PyLong_FromLong(0);
763 if (!it->index)
764 goto create_failure;
765
Guido van Rossum805365e2007-05-07 22:24:25 +0000766 return (PyObject *)it;
767
768create_failure:
Guido van Rossum317e7742007-05-08 15:18:31 +0000769 Py_DECREF(it);
Guido van Rossum805365e2007-05-07 22:24:25 +0000770 return NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000771}
772
773static PyObject *
Benjamin Petersona1864f32010-11-20 23:05:39 +0000774range_reverse(PyObject *seq)
775{
Guido van Rossum805365e2007-05-07 22:24:25 +0000776 rangeobject *range = (rangeobject*) seq;
777 longrangeiterobject *it;
778 PyObject *one, *sum, *diff, *len = NULL, *product;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000779 long lstart, lstop, lstep, new_start, new_stop;
780 unsigned long ulen;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000781
Guido van Rossum805365e2007-05-07 22:24:25 +0000782 assert(PyRange_Check(seq));
Martin v. Löwis84451042007-12-20 22:57:23 +0000783
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000784 /* reversed(range(start, stop, step)) can be expressed as
785 range(start+(n-1)*step, start-step, -step), where n is the number of
786 integers in the range.
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000787
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000788 If each of start, stop, step, -step, start-step, and the length
789 of the iterator is representable as a C long, use the int
790 version. This excludes some cases where the reversed range is
791 representable as a range_iterator, but it's good enough for
792 common cases and it makes the checks simple. */
793
794 lstart = PyLong_AsLong(range->start);
795 if (lstart == -1 && PyErr_Occurred()) {
796 PyErr_Clear();
797 goto long_range;
798 }
799 lstop = PyLong_AsLong(range->stop);
800 if (lstop == -1 && PyErr_Occurred()) {
801 PyErr_Clear();
802 goto long_range;
803 }
804 lstep = PyLong_AsLong(range->step);
805 if (lstep == -1 && PyErr_Occurred()) {
806 PyErr_Clear();
807 goto long_range;
808 }
809 /* check for possible overflow of -lstep */
810 if (lstep == LONG_MIN)
811 goto long_range;
812
813 /* check for overflow of lstart - lstep:
814
815 for lstep > 0, need only check whether lstart - lstep < LONG_MIN.
816 for lstep < 0, need only check whether lstart - lstep > LONG_MAX
817
818 Rearrange these inequalities as:
819
820 lstart - LONG_MIN < lstep (lstep > 0)
821 LONG_MAX - lstart < -lstep (lstep < 0)
822
823 and compute both sides as unsigned longs, to avoid the
824 possibility of undefined behaviour due to signed overflow. */
825
826 if (lstep > 0) {
827 if ((unsigned long)lstart - LONG_MIN < (unsigned long)lstep)
828 goto long_range;
829 }
830 else {
831 if (LONG_MAX - (unsigned long)lstart < 0UL - lstep)
832 goto long_range;
833 }
834
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000835 ulen = get_len_of_range(lstart, lstop, lstep);
836 if (ulen > (unsigned long)LONG_MAX)
837 goto long_range;
838
839 new_stop = lstart - lstep;
840 new_start = (long)(new_stop + ulen * lstep);
841 return int_range_iter(new_start, new_stop, -lstep);
842
843long_range:
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000844 it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
Guido van Rossum805365e2007-05-07 22:24:25 +0000845 if (it == NULL)
846 return NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000847
Guido van Rossum805365e2007-05-07 22:24:25 +0000848 /* start + (len - 1) * step */
849 len = range_length_obj(range);
850 if (!len)
851 goto create_failure;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000852
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000853 /* Steal reference to len. */
854 it->len = len;
855
Guido van Rossum805365e2007-05-07 22:24:25 +0000856 one = PyLong_FromLong(1);
857 if (!one)
858 goto create_failure;
859
860 diff = PyNumber_Subtract(len, one);
861 Py_DECREF(one);
862 if (!diff)
863 goto create_failure;
864
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000865 product = PyNumber_Multiply(diff, range->step);
866 Py_DECREF(diff);
Guido van Rossum805365e2007-05-07 22:24:25 +0000867 if (!product)
868 goto create_failure;
869
870 sum = PyNumber_Add(range->start, product);
871 Py_DECREF(product);
872 it->start = sum;
873 if (!it->start)
874 goto create_failure;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000875
Guido van Rossum805365e2007-05-07 22:24:25 +0000876 it->step = PyNumber_Negative(range->step);
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000877 if (!it->step)
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000878 goto create_failure;
Guido van Rossum805365e2007-05-07 22:24:25 +0000879
880 it->index = PyLong_FromLong(0);
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000881 if (!it->index)
882 goto create_failure;
Guido van Rossum805365e2007-05-07 22:24:25 +0000883
884 return (PyObject *)it;
885
886create_failure:
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000887 Py_DECREF(it);
Guido van Rossum805365e2007-05-07 22:24:25 +0000888 return NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000889}