iterobject.c 5.17 KB
Newer Older
1 2 3 4 5 6 7
/* Iterator objects */

#include "Python.h"

typedef struct {
	PyObject_HEAD
	long      it_index;
8
	PyObject *it_seq; /* Set to NULL when iterator is exhausted */
9
} seqiterobject;
10 11

PyObject *
12
PySeqIter_New(PyObject *seq)
13
{
14
	seqiterobject *it;
15 16 17 18 19

	if (!PySequence_Check(seq)) {
		PyErr_BadInternalCall();
		return NULL;
	}	
20
	it = PyObject_GC_New(seqiterobject, &PySeqIter_Type);
21 22 23 24 25
	if (it == NULL)
		return NULL;
	it->it_index = 0;
	Py_INCREF(seq);
	it->it_seq = seq;
26
	_PyObject_GC_TRACK(it);
27 28
	return (PyObject *)it;
}
29

30
static void
31
iter_dealloc(seqiterobject *it)
32
{
33
	_PyObject_GC_UNTRACK(it);
34
	Py_XDECREF(it->it_seq);
35
	PyObject_GC_Del(it);
36 37
}

38 39 40
static int
iter_traverse(seqiterobject *it, visitproc visit, void *arg)
{
41 42
	Py_VISIT(it->it_seq);
	return 0;
43 44
}

45 46 47 48 49
static PyObject *
iter_iternext(PyObject *iterator)
{
	seqiterobject *it;
	PyObject *seq;
50
	PyObject *result;
51 52 53 54

	assert(PySeqIter_Check(iterator));
	it = (seqiterobject *)iterator;
	seq = it->it_seq;
55 56
	if (seq == NULL)
		return NULL;
57

58 59 60 61 62 63 64 65 66
	result = PySequence_GetItem(seq, it->it_index);
	if (result != NULL) {
		it->it_index++;
		return result;
	}
	if (PyErr_ExceptionMatches(PyExc_IndexError) ||
	    PyErr_ExceptionMatches(PyExc_StopIteration))
	{
		PyErr_Clear();
67 68
		Py_DECREF(seq);
		it->it_seq = NULL;
69
	}
70
	return NULL;
71 72
}

73
static PyObject *
74 75
iter_len(seqiterobject *it)
{
Martin v. Löwis's avatar
Martin v. Löwis committed
76
	Py_ssize_t seqsize, len;
77 78 79 80

	if (it->it_seq) {
		seqsize = PySequence_Size(it->it_seq);
		if (seqsize == -1)
81
			return NULL;
82 83
		len = seqsize - it->it_index;
		if (len >= 0)
84
			return PyLong_FromSsize_t(len);
85
	}
86
	return PyLong_FromLong(0);
87 88
}

89
PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
90 91

static PyMethodDef seqiter_methods[] = {
92
	{"__length_hint__", (PyCFunction)iter_len, METH_NOARGS, length_hint_doc},
93
 	{NULL,		NULL}		/* sentinel */
94 95
};

96
PyTypeObject PySeqIter_Type = {
97
	PyVarObject_HEAD_INIT(&PyType_Type, 0)
98
	"iterator",				/* tp_name */
99
	sizeof(seqiterobject),			/* tp_basicsize */
100 101 102 103
	0,					/* tp_itemsize */
	/* methods */
	(destructor)iter_dealloc, 		/* tp_dealloc */
	0,					/* tp_print */
104
	0,					/* tp_getattr */
105 106 107 108
	0,					/* tp_setattr */
	0,					/* tp_compare */
	0,					/* tp_repr */
	0,					/* tp_as_number */
109
	0,					/* tp_as_sequence */
110 111 112 113
	0,					/* tp_as_mapping */
	0,					/* tp_hash */
	0,					/* tp_call */
	0,					/* tp_str */
114
	PyObject_GenericGetAttr,		/* tp_getattro */
115 116
	0,					/* tp_setattro */
	0,					/* tp_as_buffer */
117
	Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
118
 	0,					/* tp_doc */
119
 	(traverseproc)iter_traverse,		/* tp_traverse */
120 121 122
 	0,					/* tp_clear */
	0,					/* tp_richcompare */
	0,					/* tp_weaklistoffset */
123
	PyObject_SelfIter,			/* tp_iter */
124
	iter_iternext,				/* tp_iternext */
125
	seqiter_methods,			/* tp_methods */
126
	0,					/* tp_members */
127 128 129 130 131 132
};

/* -------------------------------------- */

typedef struct {
	PyObject_HEAD
133 134
	PyObject *it_callable; /* Set to NULL when iterator is exhausted */
	PyObject *it_sentinel; /* Set to NULL when iterator is exhausted */
135 136 137 138 139 140
} calliterobject;

PyObject *
PyCallIter_New(PyObject *callable, PyObject *sentinel)
{
	calliterobject *it;
141
	it = PyObject_GC_New(calliterobject, &PyCallIter_Type);
142 143 144 145 146 147
	if (it == NULL)
		return NULL;
	Py_INCREF(callable);
	it->it_callable = callable;
	Py_INCREF(sentinel);
	it->it_sentinel = sentinel;
148
	_PyObject_GC_TRACK(it);
149 150 151 152 153
	return (PyObject *)it;
}
static void
calliter_dealloc(calliterobject *it)
{
154
	_PyObject_GC_UNTRACK(it);
155 156
	Py_XDECREF(it->it_callable);
	Py_XDECREF(it->it_sentinel);
157
	PyObject_GC_Del(it);
158
}
159

160 161 162
static int
calliter_traverse(calliterobject *it, visitproc visit, void *arg)
{
163 164
	Py_VISIT(it->it_callable);
	Py_VISIT(it->it_sentinel);
165 166 167
	return 0;
}

168 169 170
static PyObject *
calliter_iternext(calliterobject *it)
{
171
	if (it->it_callable != NULL) {
172 173 174 175 176 177
		PyObject *args = PyTuple_New(0);
		PyObject *result;
		if (args == NULL)
			return NULL;
		result = PyObject_Call(it->it_callable, args, NULL);
		Py_DECREF(args);
178 179 180 181 182 183 184
		if (result != NULL) {
			int ok;
			ok = PyObject_RichCompareBool(result,
						      it->it_sentinel,
						      Py_EQ);
			if (ok == 0)
				return result; /* Common case, fast path */
185
			Py_DECREF(result);
186
			if (ok > 0) {
187 188
				Py_CLEAR(it->it_callable);
				Py_CLEAR(it->it_sentinel);
189 190 191 192
			}
		}
		else if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
			PyErr_Clear();
193 194
			Py_CLEAR(it->it_callable);
			Py_CLEAR(it->it_sentinel);
195 196
		}
	}
197
	return NULL;
198 199
}

200
PyTypeObject PyCallIter_Type = {
201
	PyVarObject_HEAD_INIT(&PyType_Type, 0)
202
	"callable_iterator",			/* tp_name */
203
	sizeof(calliterobject),			/* tp_basicsize */
204 205 206 207
	0,					/* tp_itemsize */
	/* methods */
	(destructor)calliter_dealloc, 		/* tp_dealloc */
	0,					/* tp_print */
208
	0,					/* tp_getattr */
209 210 211 212 213 214 215 216 217
	0,					/* tp_setattr */
	0,					/* tp_compare */
	0,					/* tp_repr */
	0,					/* tp_as_number */
	0,					/* tp_as_sequence */
	0,					/* tp_as_mapping */
	0,					/* tp_hash */
	0,					/* tp_call */
	0,					/* tp_str */
218
	PyObject_GenericGetAttr,		/* tp_getattro */
219 220
	0,					/* tp_setattro */
	0,					/* tp_as_buffer */
221
	Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
222 223 224
	0,					/* tp_doc */
	(traverseproc)calliter_traverse,	/* tp_traverse */
	0,					/* tp_clear */
225 226
	0,					/* tp_richcompare */
	0,					/* tp_weaklistoffset */
227
	PyObject_SelfIter,			/* tp_iter */
228
	(iternextfunc)calliter_iternext,	/* tp_iternext */
229
	0,					/* tp_methods */
230
};
231 232