blob: 965038635c39d5b1789a993c16bca591e4a6487b [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)) {
340 if (range_contains_long(r, ob))
341 Py_RETURN_TRUE;
342 else
343 Py_RETURN_FALSE;
344 } else {
345 Py_ssize_t count;
346 count = _PySequence_IterSearch((PyObject*)r, ob, PY_ITERSEARCH_COUNT);
347 if (count == -1)
348 return NULL;
349 return PyLong_FromSsize_t(count);
350 }
351}
352
353static PyObject *
354range_index(rangeobject *r, PyObject *ob)
355{
356 PyObject *idx, *tmp;
357 int contains;
358 PyObject *format_tuple, *err_string;
359 static PyObject *err_format = NULL;
360
361 if (!PyLong_CheckExact(ob) && !PyBool_Check(ob)) {
362 Py_ssize_t index;
363 index = _PySequence_IterSearch((PyObject*)r, ob, PY_ITERSEARCH_INDEX);
364 if (index == -1)
365 return NULL;
366 return PyLong_FromSsize_t(index);
367 }
368
369 contains = range_contains_long(r, ob);
370 if (contains == -1)
371 return NULL;
372
373 if (!contains)
374 goto value_error;
375
376 tmp = PyNumber_Subtract(ob, r->start);
377 if (tmp == NULL)
378 return NULL;
379
380 /* idx = (ob - r.start) // r.step */
381 idx = PyNumber_FloorDivide(tmp, r->step);
382 Py_DECREF(tmp);
383 return idx;
384
385value_error:
386
387 /* object is not in the range */
388 if (err_format == NULL) {
389 err_format = PyUnicode_FromString("%r is not in range");
390 if (err_format == NULL)
391 return NULL;
392 }
393 format_tuple = PyTuple_Pack(1, ob);
394 if (format_tuple == NULL)
395 return NULL;
396 err_string = PyUnicode_Format(err_format, format_tuple);
397 Py_DECREF(format_tuple);
398 if (err_string == NULL)
399 return NULL;
400 PyErr_SetObject(PyExc_ValueError, err_string);
401 Py_DECREF(err_string);
402 return NULL;
403}
404
Guido van Rossumc0b618a1997-05-02 03:12:38 +0000405static PySequenceMethods range_as_sequence = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000406 (lenfunc)range_length, /* sq_length */
407 0, /* sq_concat */
408 0, /* sq_repeat */
Guido van Rossum805365e2007-05-07 22:24:25 +0000409 (ssizeargfunc)range_item, /* sq_item */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000410 0, /* sq_slice */
Mark Dickinson3e124ae2009-09-22 21:47:24 +0000411 0, /* sq_ass_item */
412 0, /* sq_ass_slice */
413 (objobjproc)range_contains, /* sq_contains */
Guido van Rossum12d12c51993-10-26 17:58:25 +0000414};
415
Jeremy Hylton938ace62002-07-17 16:30:39 +0000416static PyObject * range_iter(PyObject *seq);
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000417static PyObject * range_reverse(PyObject *seq);
418
419PyDoc_STRVAR(reverse_doc,
420"Returns a reverse iterator.");
421
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000422PyDoc_STRVAR(count_doc,
423"rangeobject.count(value) -> integer -- return number of occurrences of value");
424
425PyDoc_STRVAR(index_doc,
426"rangeobject.index(value, [start, [stop]]) -> integer -- return index of value.\n"
427"Raises ValueError if the value is not present.");
428
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000429static PyMethodDef range_methods[] = {
Daniel Stutzbach9f0cbf12010-09-13 21:16:29 +0000430 {"__reversed__", (PyCFunction)range_reverse, METH_NOARGS, reverse_doc},
431 {"__reduce__", (PyCFunction)range_reduce, METH_VARARGS},
432 {"count", (PyCFunction)range_count, METH_O, count_doc},
433 {"index", (PyCFunction)range_index, METH_O, index_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000434 {NULL, NULL} /* sentinel */
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000435};
Raymond Hettinger48165d42002-06-05 20:08:48 +0000436
Guido van Rossumc0b618a1997-05-02 03:12:38 +0000437PyTypeObject PyRange_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000438 PyVarObject_HEAD_INIT(&PyType_Type, 0)
439 "range", /* Name of this type */
440 sizeof(rangeobject), /* Basic object size */
441 0, /* Item size for varobject */
442 (destructor)range_dealloc, /* tp_dealloc */
443 0, /* tp_print */
444 0, /* tp_getattr */
445 0, /* tp_setattr */
446 0, /* tp_reserved */
447 (reprfunc)range_repr, /* tp_repr */
448 0, /* tp_as_number */
449 &range_as_sequence, /* tp_as_sequence */
450 0, /* tp_as_mapping */
451 0, /* tp_hash */
452 0, /* tp_call */
453 0, /* tp_str */
454 PyObject_GenericGetAttr, /* tp_getattro */
455 0, /* tp_setattro */
456 0, /* tp_as_buffer */
457 Py_TPFLAGS_DEFAULT, /* tp_flags */
458 range_doc, /* tp_doc */
459 0, /* tp_traverse */
460 0, /* tp_clear */
461 0, /* tp_richcompare */
462 0, /* tp_weaklistoffset */
463 range_iter, /* tp_iter */
464 0, /* tp_iternext */
465 range_methods, /* tp_methods */
466 0, /* tp_members */
467 0, /* tp_getset */
468 0, /* tp_base */
469 0, /* tp_dict */
470 0, /* tp_descr_get */
471 0, /* tp_descr_set */
472 0, /* tp_dictoffset */
473 0, /* tp_init */
474 0, /* tp_alloc */
475 range_new, /* tp_new */
Guido van Rossum12d12c51993-10-26 17:58:25 +0000476};
Raymond Hettinger48165d42002-06-05 20:08:48 +0000477
Walter Dörwald4ad94212007-05-21 18:01:17 +0000478/*********************** range Iterator **************************/
Raymond Hettinger48165d42002-06-05 20:08:48 +0000479
Guido van Rossum805365e2007-05-07 22:24:25 +0000480/* There are 2 types of iterators, one for C longs, the other for
481 Python longs (ie, PyObjects). This should make iteration fast
482 in the normal case, but possible for any numeric value.
483*/
484
Raymond Hettinger48165d42002-06-05 20:08:48 +0000485typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000486 PyObject_HEAD
487 long index;
488 long start;
489 long step;
490 long len;
Raymond Hettinger48165d42002-06-05 20:08:48 +0000491} rangeiterobject;
492
Raymond Hettinger85c20a42003-11-06 14:06:48 +0000493static PyObject *
Raymond Hettinger48165d42002-06-05 20:08:48 +0000494rangeiter_next(rangeiterobject *r)
495{
Guido van Rossum805365e2007-05-07 22:24:25 +0000496 if (r->index < r->len)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000497 /* cast to unsigned to avoid possible signed overflow
Mark Dickinsonb43dbc22009-11-15 12:56:08 +0000498 in intermediate calculations. */
499 return PyLong_FromLong((long)(r->start +
500 (unsigned long)(r->index++) * r->step));
Guido van Rossum805365e2007-05-07 22:24:25 +0000501 return NULL;
Raymond Hettinger48165d42002-06-05 20:08:48 +0000502}
503
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000504static PyObject *
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000505rangeiter_len(rangeiterobject *r)
506{
Christian Heimes217cfd12007-12-02 14:31:20 +0000507 return PyLong_FromLong(r->len - r->index);
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000508}
509
Guido van Rossum805365e2007-05-07 22:24:25 +0000510typedef struct {
511 PyObject_HEAD
512 PyObject *index;
513 PyObject *start;
514 PyObject *step;
515 PyObject *len;
516} longrangeiterobject;
517
518static PyObject *
519longrangeiter_len(longrangeiterobject *r, PyObject *no_args)
520{
521 return PyNumber_Subtract(r->len, r->index);
522}
Guido van Rossum317e7742007-05-08 15:18:31 +0000523
Guido van Rossum805365e2007-05-07 22:24:25 +0000524static PyObject *rangeiter_new(PyTypeObject *, PyObject *args, PyObject *kw);
525
526PyDoc_STRVAR(length_hint_doc,
527 "Private method returning an estimate of len(list(it)).");
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000528
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000529static PyMethodDef rangeiter_methods[] = {
Guido van Rossum805365e2007-05-07 22:24:25 +0000530 {"__length_hint__", (PyCFunction)rangeiter_len, METH_NOARGS,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000531 length_hint_doc},
532 {NULL, NULL} /* sentinel */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000533};
Raymond Hettingeref9bf402004-03-10 10:10:42 +0000534
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000535PyTypeObject PyRangeIter_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000536 PyVarObject_HEAD_INIT(&PyType_Type, 0)
537 "range_iterator", /* tp_name */
538 sizeof(rangeiterobject), /* tp_basicsize */
539 0, /* tp_itemsize */
540 /* methods */
541 (destructor)PyObject_Del, /* tp_dealloc */
542 0, /* tp_print */
543 0, /* tp_getattr */
544 0, /* tp_setattr */
545 0, /* tp_reserved */
546 0, /* tp_repr */
547 0, /* tp_as_number */
548 0, /* tp_as_sequence */
549 0, /* tp_as_mapping */
550 0, /* tp_hash */
551 0, /* tp_call */
552 0, /* tp_str */
553 PyObject_GenericGetAttr, /* tp_getattro */
554 0, /* tp_setattro */
555 0, /* tp_as_buffer */
556 Py_TPFLAGS_DEFAULT, /* tp_flags */
557 0, /* tp_doc */
558 0, /* tp_traverse */
559 0, /* tp_clear */
560 0, /* tp_richcompare */
561 0, /* tp_weaklistoffset */
562 PyObject_SelfIter, /* tp_iter */
563 (iternextfunc)rangeiter_next, /* tp_iternext */
564 rangeiter_methods, /* tp_methods */
565 0, /* tp_members */
566 0, /* tp_getset */
567 0, /* tp_base */
568 0, /* tp_dict */
569 0, /* tp_descr_get */
570 0, /* tp_descr_set */
571 0, /* tp_dictoffset */
572 0, /* tp_init */
573 0, /* tp_alloc */
574 rangeiter_new, /* tp_new */
Guido van Rossum805365e2007-05-07 22:24:25 +0000575};
576
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000577/* Return number of items in range (lo, hi, step). step != 0
578 * required. The result always fits in an unsigned long.
Guido van Rossum805365e2007-05-07 22:24:25 +0000579 */
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000580static unsigned long
Guido van Rossum805365e2007-05-07 22:24:25 +0000581get_len_of_range(long lo, long hi, long step)
582{
583 /* -------------------------------------------------------------
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000584 If step > 0 and lo >= hi, or step < 0 and lo <= hi, the range is empty.
585 Else for step > 0, if n values are in the range, the last one is
Guido van Rossum805365e2007-05-07 22:24:25 +0000586 lo + (n-1)*step, which must be <= hi-1. Rearranging,
587 n <= (hi - lo - 1)/step + 1, so taking the floor of the RHS gives
588 the proper value. Since lo < hi in this case, hi-lo-1 >= 0, so
589 the RHS is non-negative and so truncation is the same as the
590 floor. Letting M be the largest positive long, the worst case
591 for the RHS numerator is hi=M, lo=-M-1, and then
592 hi-lo-1 = M-(-M-1)-1 = 2*M. Therefore unsigned long has enough
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000593 precision to compute the RHS exactly. The analysis for step < 0
594 is similar.
Guido van Rossum805365e2007-05-07 22:24:25 +0000595 ---------------------------------------------------------------*/
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000596 assert(step != 0);
597 if (step > 0 && lo < hi)
598 return 1UL + (hi - 1UL - lo) / step;
599 else if (step < 0 && lo > hi)
600 return 1UL + (lo - 1UL - hi) / (0UL - step);
601 else
602 return 0UL;
Guido van Rossum805365e2007-05-07 22:24:25 +0000603}
604
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000605/* Initialize a rangeiter object. If the length of the rangeiter object
606 is not representable as a C long, OverflowError is raised. */
607
Guido van Rossum805365e2007-05-07 22:24:25 +0000608static PyObject *
609int_range_iter(long start, long stop, long step)
610{
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000611 rangeiterobject *it = PyObject_New(rangeiterobject, &PyRangeIter_Type);
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000612 unsigned long ulen;
Guido van Rossum805365e2007-05-07 22:24:25 +0000613 if (it == NULL)
614 return NULL;
615 it->start = start;
616 it->step = step;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000617 ulen = get_len_of_range(start, stop, step);
618 if (ulen > (unsigned long)LONG_MAX) {
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000619 Py_DECREF(it);
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000620 PyErr_SetString(PyExc_OverflowError,
621 "range too large to represent as a range_iterator");
622 return NULL;
623 }
624 it->len = (long)ulen;
Guido van Rossum805365e2007-05-07 22:24:25 +0000625 it->index = 0;
626 return (PyObject *)it;
627}
628
629static PyObject *
630rangeiter_new(PyTypeObject *type, PyObject *args, PyObject *kw)
631{
632 long start, stop, step;
633
634 if (!_PyArg_NoKeywords("rangeiter()", kw))
635 return NULL;
636
637 if (!PyArg_ParseTuple(args, "lll;rangeiter() requires 3 int arguments",
638 &start, &stop, &step))
639 return NULL;
640
641 return int_range_iter(start, stop, step);
642}
643
644static PyMethodDef longrangeiter_methods[] = {
645 {"__length_hint__", (PyCFunction)longrangeiter_len, METH_NOARGS,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000646 length_hint_doc},
647 {NULL, NULL} /* sentinel */
Guido van Rossum805365e2007-05-07 22:24:25 +0000648};
649
650static void
651longrangeiter_dealloc(longrangeiterobject *r)
652{
653 Py_XDECREF(r->index);
Guido van Rossum317e7742007-05-08 15:18:31 +0000654 Py_XDECREF(r->start);
655 Py_XDECREF(r->step);
656 Py_XDECREF(r->len);
Amaury Forgeot d'Arcb7f17e42007-11-15 20:52:21 +0000657 PyObject_Del(r);
Guido van Rossum805365e2007-05-07 22:24:25 +0000658}
659
660static PyObject *
661longrangeiter_next(longrangeiterobject *r)
662{
663 PyObject *one, *product, *new_index, *result;
664 if (PyObject_RichCompareBool(r->index, r->len, Py_LT) != 1)
665 return NULL;
666
667 one = PyLong_FromLong(1);
668 if (!one)
669 return NULL;
670
Guido van Rossum805365e2007-05-07 22:24:25 +0000671 new_index = PyNumber_Add(r->index, one);
672 Py_DECREF(one);
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000673 if (!new_index)
674 return NULL;
675
676 product = PyNumber_Multiply(r->index, r->step);
677 if (!product) {
678 Py_DECREF(new_index);
Guido van Rossum805365e2007-05-07 22:24:25 +0000679 return NULL;
680 }
681
682 result = PyNumber_Add(r->start, product);
683 Py_DECREF(product);
684 if (result) {
685 Py_DECREF(r->index);
686 r->index = new_index;
687 }
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000688 else {
689 Py_DECREF(new_index);
690 }
Guido van Rossum805365e2007-05-07 22:24:25 +0000691
692 return result;
693}
694
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000695PyTypeObject PyLongRangeIter_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000696 PyVarObject_HEAD_INIT(&PyType_Type, 0)
697 "longrange_iterator", /* tp_name */
698 sizeof(longrangeiterobject), /* tp_basicsize */
699 0, /* tp_itemsize */
700 /* methods */
701 (destructor)longrangeiter_dealloc, /* tp_dealloc */
702 0, /* tp_print */
703 0, /* tp_getattr */
704 0, /* tp_setattr */
705 0, /* tp_reserved */
706 0, /* tp_repr */
707 0, /* tp_as_number */
708 0, /* tp_as_sequence */
709 0, /* tp_as_mapping */
710 0, /* tp_hash */
711 0, /* tp_call */
712 0, /* tp_str */
713 PyObject_GenericGetAttr, /* tp_getattro */
714 0, /* tp_setattro */
715 0, /* tp_as_buffer */
716 Py_TPFLAGS_DEFAULT, /* tp_flags */
717 0, /* tp_doc */
718 0, /* tp_traverse */
719 0, /* tp_clear */
720 0, /* tp_richcompare */
721 0, /* tp_weaklistoffset */
722 PyObject_SelfIter, /* tp_iter */
723 (iternextfunc)longrangeiter_next, /* tp_iternext */
724 longrangeiter_methods, /* tp_methods */
725 0,
Raymond Hettinger48165d42002-06-05 20:08:48 +0000726};
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000727
728static PyObject *
729range_iter(PyObject *seq)
730{
Guido van Rossum805365e2007-05-07 22:24:25 +0000731 rangeobject *r = (rangeobject *)seq;
732 longrangeiterobject *it;
Martin v. Löwis84451042007-12-20 22:57:23 +0000733 long lstart, lstop, lstep;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000734 PyObject *int_it;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000735
Guido van Rossum805365e2007-05-07 22:24:25 +0000736 assert(PyRange_Check(seq));
Martin v. Löwis84451042007-12-20 22:57:23 +0000737
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000738 /* If all three fields and the length convert to long, use the int
739 * version */
Martin v. Löwis84451042007-12-20 22:57:23 +0000740 lstart = PyLong_AsLong(r->start);
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000741 if (lstart == -1 && PyErr_Occurred()) {
742 PyErr_Clear();
743 goto long_range;
Martin v. Löwis84451042007-12-20 22:57:23 +0000744 }
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000745 lstop = PyLong_AsLong(r->stop);
746 if (lstop == -1 && PyErr_Occurred()) {
747 PyErr_Clear();
748 goto long_range;
749 }
750 lstep = PyLong_AsLong(r->step);
751 if (lstep == -1 && PyErr_Occurred()) {
752 PyErr_Clear();
753 goto long_range;
754 }
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000755 int_it = int_range_iter(lstart, lstop, lstep);
756 if (int_it == NULL && PyErr_ExceptionMatches(PyExc_OverflowError)) {
757 PyErr_Clear();
758 goto long_range;
759 }
760 return (PyObject *)int_it;
761
762 long_range:
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000763 it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
Guido van Rossum805365e2007-05-07 22:24:25 +0000764 if (it == NULL)
765 return NULL;
Guido van Rossum317e7742007-05-08 15:18:31 +0000766
767 /* Do all initialization here, so we can DECREF on failure. */
Guido van Rossum805365e2007-05-07 22:24:25 +0000768 it->start = r->start;
Guido van Rossum317e7742007-05-08 15:18:31 +0000769 it->step = r->step;
770 Py_INCREF(it->start);
771 Py_INCREF(it->step);
772
773 it->len = it->index = NULL;
774
Mark Dickinsoneb36d312009-06-24 18:36:46 +0000775 it->len = range_length_obj(r);
776 if (!it->len)
Guido van Rossum805365e2007-05-07 22:24:25 +0000777 goto create_failure;
Guido van Rossum805365e2007-05-07 22:24:25 +0000778 it->index = PyLong_FromLong(0);
779 if (!it->index)
780 goto create_failure;
781
Guido van Rossum805365e2007-05-07 22:24:25 +0000782 return (PyObject *)it;
783
784create_failure:
Guido van Rossum317e7742007-05-08 15:18:31 +0000785 Py_DECREF(it);
Guido van Rossum805365e2007-05-07 22:24:25 +0000786 return NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000787}
788
789static PyObject *
790range_reverse(PyObject *seq)
791{
Guido van Rossum805365e2007-05-07 22:24:25 +0000792 rangeobject *range = (rangeobject*) seq;
793 longrangeiterobject *it;
794 PyObject *one, *sum, *diff, *len = NULL, *product;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000795 long lstart, lstop, lstep, new_start, new_stop;
796 unsigned long ulen;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000797
Guido van Rossum805365e2007-05-07 22:24:25 +0000798 assert(PyRange_Check(seq));
Martin v. Löwis84451042007-12-20 22:57:23 +0000799
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000800 /* reversed(range(start, stop, step)) can be expressed as
801 range(start+(n-1)*step, start-step, -step), where n is the number of
802 integers in the range.
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000803
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000804 If each of start, stop, step, -step, start-step, and the length
805 of the iterator is representable as a C long, use the int
806 version. This excludes some cases where the reversed range is
807 representable as a range_iterator, but it's good enough for
808 common cases and it makes the checks simple. */
809
810 lstart = PyLong_AsLong(range->start);
811 if (lstart == -1 && PyErr_Occurred()) {
812 PyErr_Clear();
813 goto long_range;
814 }
815 lstop = PyLong_AsLong(range->stop);
816 if (lstop == -1 && PyErr_Occurred()) {
817 PyErr_Clear();
818 goto long_range;
819 }
820 lstep = PyLong_AsLong(range->step);
821 if (lstep == -1 && PyErr_Occurred()) {
822 PyErr_Clear();
823 goto long_range;
824 }
825 /* check for possible overflow of -lstep */
826 if (lstep == LONG_MIN)
827 goto long_range;
828
829 /* check for overflow of lstart - lstep:
830
831 for lstep > 0, need only check whether lstart - lstep < LONG_MIN.
832 for lstep < 0, need only check whether lstart - lstep > LONG_MAX
833
834 Rearrange these inequalities as:
835
836 lstart - LONG_MIN < lstep (lstep > 0)
837 LONG_MAX - lstart < -lstep (lstep < 0)
838
839 and compute both sides as unsigned longs, to avoid the
840 possibility of undefined behaviour due to signed overflow. */
841
842 if (lstep > 0) {
843 if ((unsigned long)lstart - LONG_MIN < (unsigned long)lstep)
844 goto long_range;
845 }
846 else {
847 if (LONG_MAX - (unsigned long)lstart < 0UL - lstep)
848 goto long_range;
849 }
850
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000851 ulen = get_len_of_range(lstart, lstop, lstep);
852 if (ulen > (unsigned long)LONG_MAX)
853 goto long_range;
854
855 new_stop = lstart - lstep;
856 new_start = (long)(new_stop + ulen * lstep);
857 return int_range_iter(new_start, new_stop, -lstep);
858
859long_range:
Christian Heimesa22e8bd2007-11-29 22:35:39 +0000860 it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
Guido van Rossum805365e2007-05-07 22:24:25 +0000861 if (it == NULL)
862 return NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000863
Guido van Rossum805365e2007-05-07 22:24:25 +0000864 /* start + (len - 1) * step */
865 len = range_length_obj(range);
866 if (!len)
867 goto create_failure;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000868
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000869 /* Steal reference to len. */
870 it->len = len;
871
Guido van Rossum805365e2007-05-07 22:24:25 +0000872 one = PyLong_FromLong(1);
873 if (!one)
874 goto create_failure;
875
876 diff = PyNumber_Subtract(len, one);
877 Py_DECREF(one);
878 if (!diff)
879 goto create_failure;
880
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000881 product = PyNumber_Multiply(diff, range->step);
882 Py_DECREF(diff);
Guido van Rossum805365e2007-05-07 22:24:25 +0000883 if (!product)
884 goto create_failure;
885
886 sum = PyNumber_Add(range->start, product);
887 Py_DECREF(product);
888 it->start = sum;
889 if (!it->start)
890 goto create_failure;
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000891
Guido van Rossum805365e2007-05-07 22:24:25 +0000892 it->step = PyNumber_Negative(range->step);
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000893 if (!it->step)
Mark Dickinsond550c9a2009-11-15 09:57:26 +0000894 goto create_failure;
Guido van Rossum805365e2007-05-07 22:24:25 +0000895
896 it->index = PyLong_FromLong(0);
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000897 if (!it->index)
898 goto create_failure;
Guido van Rossum805365e2007-05-07 22:24:25 +0000899
900 return (PyObject *)it;
901
902create_failure:
Benjamin Peterson36fbb732009-11-16 00:34:25 +0000903 Py_DECREF(it);
Guido van Rossum805365e2007-05-07 22:24:25 +0000904 return NULL;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000905}