blob: d7bf018a1d40be0b98379041158e7b4e7be9accf [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 *
23validate_step(PyObject *step)
Raymond Hettingerc4c453f2002-06-05 23:12:45 +000024{
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 *
52range_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
121range_dealloc(rangeobject *r)
Guido van Rossum12d12c51993-10-26 17:58:25 +0000122{
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*
134range_length_obj(rangeobject *r)
135{
136 /* -------------------------------------------------------------
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
Fred Drake45cfbcc2000-07-09 06:21:27 +0000207range_length(rangeobject *r)
Guido van Rossum12d12c51993-10-26 17:58:25 +0000208{
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 *
221range_item(rangeobject *r, Py_ssize_t i)
222{
223 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 *
Fred Drake45cfbcc2000-07-09 06:21:27 +0000249range_repr(rangeobject *r)
Guido van Rossum12d12c51993-10-26 17:58:25 +0000250{
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 *
270range_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
278range_contains_long(rangeobject *r, PyObject *ob)
279{
280 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
328range_contains(rangeobject *r, PyObject *ob) {
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000329 if (PyLong_CheckExact(ob) || PyBool_Check(ob))
330 return range_contains_long(r, ob);
Mark Dickinson3e124ae2009-09-22 21:47:24 +0000331
Mark Dickinson3e124ae2009-09-22 21:47:24 +0000332 return (int)_PySequence_IterSearch((PyObject*)r, ob,
333 PY_ITERSEARCH_CONTAINS);
334}
335
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000336static PyObject *
337range_count(rangeobject *r, PyObject *ob)
338{
339 if (PyLong_CheckExact(ob) || PyBool_Check(ob)) {
Georg Brandl7e5343b2010-11-20 22:40:10 +0000340 int result = range_contains_long(r, ob);
341 if (result == -1)
342 return NULL;
343 else if (result)
Benjamin Peterson0b458d52010-11-20 22:35:41 +0000344 return PyLong_FromLong(1);
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000345 else
Benjamin Peterson0b458d52010-11-20 22:35:41 +0000346 return PyLong_FromLong(0);
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000347 } else {
348 Py_ssize_t count;
349 count = _PySequence_IterSearch((PyObject*)r, ob, PY_ITERSEARCH_COUNT);
350 if (count == -1)
351 return NULL;
352 return PyLong_FromSsize_t(count);
353 }
354}
355
356static PyObject *
357range_index(rangeobject *r, PyObject *ob)
358{
359 PyObject *idx, *tmp;
360 int contains;
361 PyObject *format_tuple, *err_string;
362 static PyObject *err_format = NULL;
363
364 if (!PyLong_CheckExact(ob) && !PyBool_Check(ob)) {
365 Py_ssize_t index;
366 index = _PySequence_IterSearch((PyObject*)r, ob, PY_ITERSEARCH_INDEX);
367 if (index == -1)
368 return NULL;
369 return PyLong_FromSsize_t(index);
370 }
371
372 contains = range_contains_long(r, ob);
373 if (contains == -1)
374 return NULL;
375
376 if (!contains)
377 goto value_error;
378
379 tmp = PyNumber_Subtract(ob, r->start);
380 if (tmp == NULL)
381 return NULL;
382
383 /* idx = (ob - r.start) // r.step */
384 idx = PyNumber_FloorDivide(tmp, r->step);
385 Py_DECREF(tmp);
386 return idx;
387
388value_error:
389
390 /* object is not in the range */
391 if (err_format == NULL) {
392 err_format = PyUnicode_FromString("%r is not in range");
393 if (err_format == NULL)
394 return NULL;
395 }
396 format_tuple = PyTuple_Pack(1, ob);
397 if (format_tuple == NULL)
398 return NULL;
399 err_string = PyUnicode_Format(err_format, format_tuple);
400 Py_DECREF(format_tuple);
401 if (err_string == NULL)
402 return NULL;
403 PyErr_SetObject(PyExc_ValueError, err_string);
404 Py_DECREF(err_string);
405 return NULL;
406}
407
Guido van Rossumc0b618a1997-05-02 03:12:38 +0000408static PySequenceMethods range_as_sequence = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000409 (lenfunc)range_length, /* sq_length */
410 0, /* sq_concat */
411 0, /* sq_repeat */
Guido van Rossum805365e2007-05-07 22:24:25 +0000412 (ssizeargfunc)range_item, /* sq_item */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000413 0, /* sq_slice */
Mark Dickinson3e124ae2009-09-22 21:47:24 +0000414 0, /* sq_ass_item */
415 0, /* sq_ass_slice */
416 (objobjproc)range_contains, /* sq_contains */
Guido van Rossum12d12c51993-10-26 17:58:25 +0000417};
418
Jeremy Hylton938ace62002-07-17 16:30:39 +0000419static PyObject * range_iter(PyObject *seq);
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000420static PyObject * range_reverse(PyObject *seq);
421
422PyDoc_STRVAR(reverse_doc,
423"Returns a reverse iterator.");
424
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000425PyDoc_STRVAR(count_doc,
426"rangeobject.count(value) -> integer -- return number of occurrences of value");
427
428PyDoc_STRVAR(index_doc,
429"rangeobject.index(value, [start, [stop]]) -> integer -- return index of value.\n"
430"Raises ValueError if the value is not present.");
431
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000432static PyMethodDef range_methods[] = {
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000433 {"__reversed__", (PyCFunction)range_reverse, METH_NOARGS, reverse_doc},
434 {"__reduce__", (PyCFunction)range_reduce, METH_VARARGS},
435 {"count", (PyCFunction)range_count, METH_O, count_doc},
436 {"index", (PyCFunction)range_index, METH_O, index_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000437 {NULL, NULL} /* sentinel */
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000438};
Raymond Hettinger48165d42002-06-05 20:08:48 +0000439
Guido van Rossumc0b618a1997-05-02 03:12:38 +0000440PyTypeObject PyRange_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000441 PyVarObject_HEAD_INIT(&PyType_Type, 0)
442 "range", /* Name of this type */
443 sizeof(rangeobject), /* Basic object size */
444 0, /* Item size for varobject */
445 (destructor)range_dealloc, /* tp_dealloc */
446 0, /* tp_print */
447 0, /* tp_getattr */
448 0, /* tp_setattr */
449 0, /* tp_reserved */
450 (reprfunc)range_repr, /* tp_repr */
451 0, /* tp_as_number */
452 &range_as_sequence, /* tp_as_sequence */
453 0, /* tp_as_mapping */
454 0, /* tp_hash */
455 0, /* tp_call */
456 0, /* tp_str */
457 PyObject_GenericGetAttr, /* tp_getattro */
458 0, /* tp_setattro */
459 0, /* tp_as_buffer */
460 Py_TPFLAGS_DEFAULT, /* tp_flags */
461 range_doc, /* tp_doc */
462 0, /* tp_traverse */
463 0, /* tp_clear */
464 0, /* tp_richcompare */
465 0, /* tp_weaklistoffset */
466 range_iter, /* tp_iter */
467 0, /* tp_iternext */
468 range_methods, /* tp_methods */
469 0, /* tp_members */
470 0, /* tp_getset */
471 0, /* tp_base */
472 0, /* tp_dict */
473 0, /* tp_descr_get */
474 0, /* tp_descr_set */
475 0, /* tp_dictoffset */
476 0, /* tp_init */
477 0, /* tp_alloc */
478 range_new, /* tp_new */
Guido van Rossum12d12c51993-10-26 17:58:25 +0000479};
Raymond Hettinger48165d42002-06-05 20:08:48 +0000480
Walter Dörwald4ad94212007-05-21 18:01:17 +0000481/*********************** range Iterator **************************/
Raymond Hettinger48165d42002-06-05 20:08:48 +0000482
Guido van Rossum805365e2007-05-07 22:24:25 +0000483/* There are 2 types of iterators, one for C longs, the other for
484 Python longs (ie, PyObjects). This should make iteration fast
485 in the normal case, but possible for any numeric value.
486*/
487
Raymond Hettinger48165d42002-06-05 20:08:48 +0000488typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000489 PyObject_HEAD
490 long index;
491 long start;
492 long step;
493 long len;
Raymond Hettinger48165d42002-06-05 20:08:48 +0000494} rangeiterobject;
495
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000496static PyObject *
Raymond Hettinger48165d42002-06-05 20:08:48 +0000497rangeiter_next(rangeiterobject *r)
498{
Guido van Rossum805365e2007-05-07 22:24:25 +0000499 if (r->index < r->len)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000500 /* cast to unsigned to avoid possible signed overflow
Mark Dickinsonb43dbc22009-11-15 12:56:08 +0000501 in intermediate calculations. */
502 return PyLong_FromLong((long)(r->start +
503 (unsigned long)(r->index++) * r->step));
Guido van Rossum805365e2007-05-07 22:24:25 +0000504 return NULL;
Raymond Hettinger48165d42002-06-05 20:08:48 +0000505}
506
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000507static PyObject *
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000508rangeiter_len(rangeiterobject *r)
509{
Christian Heimes217cfd12007-12-02 14:31:20 +0000510 return PyLong_FromLong(r->len - r->index);
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000511}
512
Guido van Rossum805365e2007-05-07 22:24:25 +0000513typedef struct {
514 PyObject_HEAD
515 PyObject *index;
516 PyObject *start;
517 PyObject *step;
518 PyObject *len;
519} longrangeiterobject;
520
521static PyObject *
522longrangeiter_len(longrangeiterobject *r, PyObject *no_args)
523{
524 return PyNumber_Subtract(r->len, r->index);
525}
Guido van Rossum317e7742007-05-08 15:18:31 +0000526
Guido van Rossum805365e2007-05-07 22:24:25 +0000527static PyObject *rangeiter_new(PyTypeObject *, PyObject *args, PyObject *kw);
528
529PyDoc_STRVAR(length_hint_doc,
530 "Private method returning an estimate of len(list(it)).");
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000531
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000532static PyMethodDef rangeiter_methods[] = {
Guido van Rossum805365e2007-05-07 22:24:25 +0000533 {"__length_hint__", (PyCFunction)rangeiter_len, METH_NOARGS,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000534 length_hint_doc},
535 {NULL, NULL} /* sentinel */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000536};
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000537
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000538PyTypeObject PyRangeIter_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000539 PyVarObject_HEAD_INIT(&PyType_Type, 0)
540 "range_iterator", /* tp_name */
541 sizeof(rangeiterobject), /* tp_basicsize */
542 0, /* tp_itemsize */
543 /* methods */
544 (destructor)PyObject_Del, /* tp_dealloc */
545 0, /* tp_print */
546 0, /* tp_getattr */
547 0, /* tp_setattr */
548 0, /* tp_reserved */
549 0, /* tp_repr */
550 0, /* tp_as_number */
551 0, /* tp_as_sequence */
552 0, /* tp_as_mapping */
553 0, /* tp_hash */
554 0, /* tp_call */
555 0, /* tp_str */
556 PyObject_GenericGetAttr, /* tp_getattro */
557 0, /* tp_setattro */
558 0, /* tp_as_buffer */
559 Py_TPFLAGS_DEFAULT, /* tp_flags */
560 0, /* tp_doc */
561 0, /* tp_traverse */
562 0, /* tp_clear */
563 0, /* tp_richcompare */
564 0, /* tp_weaklistoffset */
565 PyObject_SelfIter, /* tp_iter */
566 (iternextfunc)rangeiter_next, /* tp_iternext */
567 rangeiter_methods, /* tp_methods */
568 0, /* tp_members */
569 0, /* tp_getset */
570 0, /* tp_base */
571 0, /* tp_dict */
572 0, /* tp_descr_get */
573 0, /* tp_descr_set */
574 0, /* tp_dictoffset */
575 0, /* tp_init */
576 0, /* tp_alloc */
577 rangeiter_new, /* tp_new */
Guido van Rossum805365e2007-05-07 22:24:25 +0000578};
579
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000580/* Return number of items in range (lo, hi, step). step != 0
581 * required. The result always fits in an unsigned long.
Guido van Rossum805365e2007-05-07 22:24:25 +0000582 */
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000583static unsigned long
Guido van Rossum805365e2007-05-07 22:24:25 +0000584get_len_of_range(long lo, long hi, long step)
585{
586 /* -------------------------------------------------------------
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000587 If step > 0 and lo >= hi, or step < 0 and lo <= hi, the range is empty.
588 Else for step > 0, if n values are in the range, the last one is
Guido van Rossum805365e2007-05-07 22:24:25 +0000589 lo + (n-1)*step, which must be <= hi-1. Rearranging,
590 n <= (hi - lo - 1)/step + 1, so taking the floor of the RHS gives
591 the proper value. Since lo < hi in this case, hi-lo-1 >= 0, so
592 the RHS is non-negative and so truncation is the same as the
593 floor. Letting M be the largest positive long, the worst case
594 for the RHS numerator is hi=M, lo=-M-1, and then
595 hi-lo-1 = M-(-M-1)-1 = 2*M. Therefore unsigned long has enough
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000596 precision to compute the RHS exactly. The analysis for step < 0
597 is similar.
Guido van Rossum805365e2007-05-07 22:24:25 +0000598 ---------------------------------------------------------------*/
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000599 assert(step != 0);
600 if (step > 0 && lo < hi)
601 return 1UL + (hi - 1UL - lo) / step;
602 else if (step < 0 && lo > hi)
603 return 1UL + (lo - 1UL - hi) / (0UL - step);
604 else
605 return 0UL;
Guido van Rossum805365e2007-05-07 22:24:25 +0000606}
607
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000608/* Initialize a rangeiter object. If the length of the rangeiter object
609 is not representable as a C long, OverflowError is raised. */
610
Guido van Rossum805365e2007-05-07 22:24:25 +0000611static PyObject *
612int_range_iter(long start, long stop, long step)
613{
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000614 rangeiterobject *it = PyObject_New(rangeiterobject, &PyRangeIter_Type);
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000615 unsigned long ulen;
Guido van Rossum805365e2007-05-07 22:24:25 +0000616 if (it == NULL)
617 return NULL;
618 it->start = start;
619 it->step = step;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000620 ulen = get_len_of_range(start, stop, step);
621 if (ulen > (unsigned long)LONG_MAX) {
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000622 Py_DECREF(it);
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000623 PyErr_SetString(PyExc_OverflowError,
624 "range too large to represent as a range_iterator");
625 return NULL;
626 }
627 it->len = (long)ulen;
Guido van Rossum805365e2007-05-07 22:24:25 +0000628 it->index = 0;
629 return (PyObject *)it;
630}
631
632static PyObject *
633rangeiter_new(PyTypeObject *type, PyObject *args, PyObject *kw)
634{
635 long start, stop, step;
636
637 if (!_PyArg_NoKeywords("rangeiter()", kw))
638 return NULL;
639
640 if (!PyArg_ParseTuple(args, "lll;rangeiter() requires 3 int arguments",
641 &start, &stop, &step))
642 return NULL;
643
644 return int_range_iter(start, stop, step);
645}
646
647static PyMethodDef longrangeiter_methods[] = {
648 {"__length_hint__", (PyCFunction)longrangeiter_len, METH_NOARGS,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000649 length_hint_doc},
650 {NULL, NULL} /* sentinel */
Guido van Rossum805365e2007-05-07 22:24:25 +0000651};
652
653static void
654longrangeiter_dealloc(longrangeiterobject *r)
655{
656 Py_XDECREF(r->index);
Guido van Rossum317e7742007-05-08 15:18:31 +0000657 Py_XDECREF(r->start);
658 Py_XDECREF(r->step);
659 Py_XDECREF(r->len);
Amaury Forgeot d'Arcb7f17e42007-11-15 20:52:21 +0000660 PyObject_Del(r);
Guido van Rossum805365e2007-05-07 22:24:25 +0000661}
662
663static PyObject *
664longrangeiter_next(longrangeiterobject *r)
665{
666 PyObject *one, *product, *new_index, *result;
667 if (PyObject_RichCompareBool(r->index, r->len, Py_LT) != 1)
668 return NULL;
669
670 one = PyLong_FromLong(1);
671 if (!one)
672 return NULL;
673
Guido van Rossum805365e2007-05-07 22:24:25 +0000674 new_index = PyNumber_Add(r->index, one);
675 Py_DECREF(one);
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000676 if (!new_index)
677 return NULL;
678
679 product = PyNumber_Multiply(r->index, r->step);
680 if (!product) {
681 Py_DECREF(new_index);
Guido van Rossum805365e2007-05-07 22:24:25 +0000682 return NULL;
683 }
684
685 result = PyNumber_Add(r->start, product);
686 Py_DECREF(product);
687 if (result) {
688 Py_DECREF(r->index);
689 r->index = new_index;
690 }
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000691 else {
692 Py_DECREF(new_index);
693 }
Guido van Rossum805365e2007-05-07 22:24:25 +0000694
695 return result;
696}
697
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000698PyTypeObject PyLongRangeIter_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000699 PyVarObject_HEAD_INIT(&PyType_Type, 0)
700 "longrange_iterator", /* tp_name */
701 sizeof(longrangeiterobject), /* tp_basicsize */
702 0, /* tp_itemsize */
703 /* methods */
704 (destructor)longrangeiter_dealloc, /* tp_dealloc */
705 0, /* tp_print */
706 0, /* tp_getattr */
707 0, /* tp_setattr */
708 0, /* tp_reserved */
709 0, /* tp_repr */
710 0, /* tp_as_number */
711 0, /* tp_as_sequence */
712 0, /* tp_as_mapping */
713 0, /* tp_hash */
714 0, /* tp_call */
715 0, /* tp_str */
716 PyObject_GenericGetAttr, /* tp_getattro */
717 0, /* tp_setattro */
718 0, /* tp_as_buffer */
719 Py_TPFLAGS_DEFAULT, /* tp_flags */
720 0, /* tp_doc */
721 0, /* tp_traverse */
722 0, /* tp_clear */
723 0, /* tp_richcompare */
724 0, /* tp_weaklistoffset */
725 PyObject_SelfIter, /* tp_iter */
726 (iternextfunc)longrangeiter_next, /* tp_iternext */
727 longrangeiter_methods, /* tp_methods */
728 0,
Raymond Hettinger48165d42002-06-05 20:08:48 +0000729};
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000730
731static PyObject *
732range_iter(PyObject *seq)
733{
Guido van Rossum805365e2007-05-07 22:24:25 +0000734 rangeobject *r = (rangeobject *)seq;
735 longrangeiterobject *it;
Martin v. Löwis84451042007-12-20 22:57:23 +0000736 long lstart, lstop, lstep;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000737 PyObject *int_it;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000738
Guido van Rossum805365e2007-05-07 22:24:25 +0000739 assert(PyRange_Check(seq));
Martin v. Löwis84451042007-12-20 22:57:23 +0000740
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000741 /* If all three fields and the length convert to long, use the int
742 * version */
Martin v. Löwis84451042007-12-20 22:57:23 +0000743 lstart = PyLong_AsLong(r->start);
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000744 if (lstart == -1 && PyErr_Occurred()) {
745 PyErr_Clear();
746 goto long_range;
Martin v. Löwis84451042007-12-20 22:57:23 +0000747 }
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000748 lstop = PyLong_AsLong(r->stop);
749 if (lstop == -1 && PyErr_Occurred()) {
750 PyErr_Clear();
751 goto long_range;
752 }
753 lstep = PyLong_AsLong(r->step);
754 if (lstep == -1 && PyErr_Occurred()) {
755 PyErr_Clear();
756 goto long_range;
757 }
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000758 int_it = int_range_iter(lstart, lstop, lstep);
759 if (int_it == NULL && PyErr_ExceptionMatches(PyExc_OverflowError)) {
760 PyErr_Clear();
761 goto long_range;
762 }
763 return (PyObject *)int_it;
764
765 long_range:
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000766 it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
Guido van Rossum805365e2007-05-07 22:24:25 +0000767 if (it == NULL)
768 return NULL;
Guido van Rossum317e7742007-05-08 15:18:31 +0000769
770 /* Do all initialization here, so we can DECREF on failure. */
Guido van Rossum805365e2007-05-07 22:24:25 +0000771 it->start = r->start;
Guido van Rossum317e7742007-05-08 15:18:31 +0000772 it->step = r->step;
773 Py_INCREF(it->start);
774 Py_INCREF(it->step);
775
776 it->len = it->index = NULL;
777
Mark Dickinsoneb36d312009-06-24 18:36:46 +0000778 it->len = range_length_obj(r);
779 if (!it->len)
Guido van Rossum805365e2007-05-07 22:24:25 +0000780 goto create_failure;
Guido van Rossum805365e2007-05-07 22:24:25 +0000781 it->index = PyLong_FromLong(0);
782 if (!it->index)
783 goto create_failure;
784
Guido van Rossum805365e2007-05-07 22:24:25 +0000785 return (PyObject *)it;
786
787create_failure:
Guido van Rossum317e7742007-05-08 15:18:31 +0000788 Py_DECREF(it);
Guido van Rossum805365e2007-05-07 22:24:25 +0000789 return NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000790}
791
792static PyObject *
793range_reverse(PyObject *seq)
794{
Guido van Rossum805365e2007-05-07 22:24:25 +0000795 rangeobject *range = (rangeobject*) seq;
796 longrangeiterobject *it;
797 PyObject *one, *sum, *diff, *len = NULL, *product;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000798 long lstart, lstop, lstep, new_start, new_stop;
799 unsigned long ulen;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000800
Guido van Rossum805365e2007-05-07 22:24:25 +0000801 assert(PyRange_Check(seq));
Martin v. Löwis84451042007-12-20 22:57:23 +0000802
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000803 /* reversed(range(start, stop, step)) can be expressed as
804 range(start+(n-1)*step, start-step, -step), where n is the number of
805 integers in the range.
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000806
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000807 If each of start, stop, step, -step, start-step, and the length
808 of the iterator is representable as a C long, use the int
809 version. This excludes some cases where the reversed range is
810 representable as a range_iterator, but it's good enough for
811 common cases and it makes the checks simple. */
812
813 lstart = PyLong_AsLong(range->start);
814 if (lstart == -1 && PyErr_Occurred()) {
815 PyErr_Clear();
816 goto long_range;
817 }
818 lstop = PyLong_AsLong(range->stop);
819 if (lstop == -1 && PyErr_Occurred()) {
820 PyErr_Clear();
821 goto long_range;
822 }
823 lstep = PyLong_AsLong(range->step);
824 if (lstep == -1 && PyErr_Occurred()) {
825 PyErr_Clear();
826 goto long_range;
827 }
828 /* check for possible overflow of -lstep */
829 if (lstep == LONG_MIN)
830 goto long_range;
831
832 /* check for overflow of lstart - lstep:
833
834 for lstep > 0, need only check whether lstart - lstep < LONG_MIN.
835 for lstep < 0, need only check whether lstart - lstep > LONG_MAX
836
837 Rearrange these inequalities as:
838
839 lstart - LONG_MIN < lstep (lstep > 0)
840 LONG_MAX - lstart < -lstep (lstep < 0)
841
842 and compute both sides as unsigned longs, to avoid the
843 possibility of undefined behaviour due to signed overflow. */
844
845 if (lstep > 0) {
846 if ((unsigned long)lstart - LONG_MIN < (unsigned long)lstep)
847 goto long_range;
848 }
849 else {
850 if (LONG_MAX - (unsigned long)lstart < 0UL - lstep)
851 goto long_range;
852 }
853
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000854 ulen = get_len_of_range(lstart, lstop, lstep);
855 if (ulen > (unsigned long)LONG_MAX)
856 goto long_range;
857
858 new_stop = lstart - lstep;
859 new_start = (long)(new_stop + ulen * lstep);
860 return int_range_iter(new_start, new_stop, -lstep);
861
862long_range:
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000863 it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
Guido van Rossum805365e2007-05-07 22:24:25 +0000864 if (it == NULL)
865 return NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000866
Guido van Rossum805365e2007-05-07 22:24:25 +0000867 /* start + (len - 1) * step */
868 len = range_length_obj(range);
869 if (!len)
870 goto create_failure;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000871
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000872 /* Steal reference to len. */
873 it->len = len;
874
Guido van Rossum805365e2007-05-07 22:24:25 +0000875 one = PyLong_FromLong(1);
876 if (!one)
877 goto create_failure;
878
879 diff = PyNumber_Subtract(len, one);
880 Py_DECREF(one);
881 if (!diff)
882 goto create_failure;
883
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000884 product = PyNumber_Multiply(diff, range->step);
885 Py_DECREF(diff);
Guido van Rossum805365e2007-05-07 22:24:25 +0000886 if (!product)
887 goto create_failure;
888
889 sum = PyNumber_Add(range->start, product);
890 Py_DECREF(product);
891 it->start = sum;
892 if (!it->start)
893 goto create_failure;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000894
Guido van Rossum805365e2007-05-07 22:24:25 +0000895 it->step = PyNumber_Negative(range->step);
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000896 if (!it->step)
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000897 goto create_failure;
Guido van Rossum805365e2007-05-07 22:24:25 +0000898
899 it->index = PyLong_FromLong(0);
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000900 if (!it->index)
901 goto create_failure;
Guido van Rossum805365e2007-05-07 22:24:25 +0000902
903 return (PyObject *)it;
904
905create_failure:
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000906 Py_DECREF(it);
Guido van Rossum805365e2007-05-07 22:24:25 +0000907 return NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000908}