Guido van Rossum | 7dab242 | 2002-04-26 19:40:56 +0000 | [diff] [blame] | 1 | /* enumerate object */ |
| 2 | |
| 3 | #include "Python.h" |
| 4 | |
| 5 | typedef struct { |
| 6 | PyObject_HEAD |
| 7 | long en_index; /* current index of enumeration */ |
| 8 | PyObject* en_sit; /* secondary iterator of enumeration */ |
Raymond Hettinger | e8b0f04 | 2003-05-28 14:05:34 +0000 | [diff] [blame] | 9 | PyObject* en_result; /* result tuple */ |
Guido van Rossum | 7dab242 | 2002-04-26 19:40:56 +0000 | [diff] [blame] | 10 | } enumobject; |
| 11 | |
| 12 | PyTypeObject PyEnum_Type; |
| 13 | |
| 14 | static PyObject * |
| 15 | enum_new(PyTypeObject *type, PyObject *args, PyObject *kwds) |
| 16 | { |
| 17 | enumobject *en; |
| 18 | PyObject *seq = NULL; |
| 19 | static char *kwlist[] = {"sequence", 0}; |
| 20 | |
| 21 | if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:enumerate", kwlist, |
| 22 | &seq)) |
| 23 | return NULL; |
| 24 | |
| 25 | en = (enumobject *)type->tp_alloc(type, 0); |
| 26 | if (en == NULL) |
| 27 | return NULL; |
| 28 | en->en_index = 0; |
| 29 | en->en_sit = PyObject_GetIter(seq); |
| 30 | if (en->en_sit == NULL) { |
| 31 | Py_DECREF(en); |
| 32 | return NULL; |
| 33 | } |
Raymond Hettinger | e8b0f04 | 2003-05-28 14:05:34 +0000 | [diff] [blame] | 34 | en->en_result = PyTuple_New(2); |
| 35 | if (en->en_result == NULL) { |
| 36 | Py_DECREF(en->en_sit); |
| 37 | Py_DECREF(en); |
| 38 | return NULL; |
| 39 | } |
| 40 | Py_INCREF(Py_None); |
| 41 | PyTuple_SET_ITEM(en->en_result, 0, Py_None); |
| 42 | Py_INCREF(Py_None); |
| 43 | PyTuple_SET_ITEM(en->en_result, 1, Py_None); |
Guido van Rossum | 7dab242 | 2002-04-26 19:40:56 +0000 | [diff] [blame] | 44 | return (PyObject *)en; |
| 45 | } |
| 46 | |
| 47 | static void |
| 48 | enum_dealloc(enumobject *en) |
| 49 | { |
| 50 | PyObject_GC_UnTrack(en); |
| 51 | Py_XDECREF(en->en_sit); |
Raymond Hettinger | e8b0f04 | 2003-05-28 14:05:34 +0000 | [diff] [blame] | 52 | Py_XDECREF(en->en_result); |
Guido van Rossum | 7dab242 | 2002-04-26 19:40:56 +0000 | [diff] [blame] | 53 | en->ob_type->tp_free(en); |
| 54 | } |
| 55 | |
| 56 | static int |
| 57 | enum_traverse(enumobject *en, visitproc visit, void *arg) |
| 58 | { |
Raymond Hettinger | e8b0f04 | 2003-05-28 14:05:34 +0000 | [diff] [blame] | 59 | int err; |
| 60 | |
| 61 | if (en->en_sit) { |
| 62 | err = visit(en->en_sit, arg); |
| 63 | if (err) |
| 64 | return err; |
| 65 | } |
| 66 | if (en->en_result) { |
| 67 | err = visit(en->en_result, arg); |
| 68 | if (err) |
| 69 | return err; |
| 70 | } |
Guido van Rossum | 7dab242 | 2002-04-26 19:40:56 +0000 | [diff] [blame] | 71 | return 0; |
| 72 | } |
| 73 | |
| 74 | static PyObject * |
| 75 | enum_next(enumobject *en) |
| 76 | { |
Guido van Rossum | 7dab242 | 2002-04-26 19:40:56 +0000 | [diff] [blame] | 77 | PyObject *next_index; |
Guido van Rossum | ca5ed5b | 2002-07-16 21:02:42 +0000 | [diff] [blame] | 78 | PyObject *next_item; |
Raymond Hettinger | e8b0f04 | 2003-05-28 14:05:34 +0000 | [diff] [blame] | 79 | PyObject *result = en->en_result; |
| 80 | PyObject *it = en->en_sit; |
Guido van Rossum | 7dab242 | 2002-04-26 19:40:56 +0000 | [diff] [blame] | 81 | |
Raymond Hettinger | e8b0f04 | 2003-05-28 14:05:34 +0000 | [diff] [blame] | 82 | next_item = (*it->ob_type->tp_iternext)(it); |
| 83 | if (next_item == NULL) |
Guido van Rossum | ca5ed5b | 2002-07-16 21:02:42 +0000 | [diff] [blame] | 84 | return NULL; |
| 85 | |
| 86 | next_index = PyInt_FromLong(en->en_index); |
| 87 | if (next_index == NULL) { |
Raymond Hettinger | e8b0f04 | 2003-05-28 14:05:34 +0000 | [diff] [blame] | 88 | Py_DECREF(next_item); |
Guido van Rossum | 7dab242 | 2002-04-26 19:40:56 +0000 | [diff] [blame] | 89 | return NULL; |
| 90 | } |
Raymond Hettinger | e8b0f04 | 2003-05-28 14:05:34 +0000 | [diff] [blame] | 91 | en->en_index++; |
| 92 | |
| 93 | if (result->ob_refcnt == 1) { |
| 94 | Py_INCREF(result); |
| 95 | Py_DECREF(PyTuple_GET_ITEM(result, 0)); |
| 96 | Py_DECREF(PyTuple_GET_ITEM(result, 1)); |
| 97 | } else { |
| 98 | result = PyTuple_New(2); |
| 99 | if (result == NULL) { |
| 100 | Py_DECREF(next_index); |
| 101 | Py_DECREF(next_item); |
| 102 | return NULL; |
| 103 | } |
| 104 | } |
Guido van Rossum | ca5ed5b | 2002-07-16 21:02:42 +0000 | [diff] [blame] | 105 | PyTuple_SET_ITEM(result, 0, next_index); |
Guido van Rossum | 7dab242 | 2002-04-26 19:40:56 +0000 | [diff] [blame] | 106 | PyTuple_SET_ITEM(result, 1, next_item); |
| 107 | return result; |
| 108 | } |
| 109 | |
Martin v. Löwis | 14f8b4c | 2002-06-13 20:33:02 +0000 | [diff] [blame] | 110 | PyDoc_STRVAR(enum_doc, |
Jeremy Hylton | fbbe347 | 2003-04-21 20:26:25 +0000 | [diff] [blame] | 111 | "enumerate(iterable) -> iterator for index, value of iterable\n" |
| 112 | "\n" |
| 113 | "Return an enumerate object. iterable must be an other object that supports\n" |
| 114 | "iteration. The enumerate object yields pairs containing a count (from\n" |
| 115 | "zero) and a value yielded by the iterable argument. enumerate is useful\n" |
| 116 | "for obtaining an indexed list: (0, seq[0]), (1, seq[1]), (2, seq[2]), ..."); |
Guido van Rossum | 7dab242 | 2002-04-26 19:40:56 +0000 | [diff] [blame] | 117 | |
| 118 | PyTypeObject PyEnum_Type = { |
| 119 | PyObject_HEAD_INIT(&PyType_Type) |
| 120 | 0, /* ob_size */ |
| 121 | "enumerate", /* tp_name */ |
| 122 | sizeof(enumobject), /* tp_basicsize */ |
| 123 | 0, /* tp_itemsize */ |
| 124 | /* methods */ |
| 125 | (destructor)enum_dealloc, /* tp_dealloc */ |
| 126 | 0, /* tp_print */ |
| 127 | 0, /* tp_getattr */ |
| 128 | 0, /* tp_setattr */ |
| 129 | 0, /* tp_compare */ |
| 130 | 0, /* tp_repr */ |
| 131 | 0, /* tp_as_number */ |
| 132 | 0, /* tp_as_sequence */ |
| 133 | 0, /* tp_as_mapping */ |
| 134 | 0, /* tp_hash */ |
| 135 | 0, /* tp_call */ |
| 136 | 0, /* tp_str */ |
| 137 | PyObject_GenericGetAttr, /* tp_getattro */ |
| 138 | 0, /* tp_setattro */ |
| 139 | 0, /* tp_as_buffer */ |
| 140 | Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | |
| 141 | Py_TPFLAGS_BASETYPE, /* tp_flags */ |
| 142 | enum_doc, /* tp_doc */ |
| 143 | (traverseproc)enum_traverse, /* tp_traverse */ |
| 144 | 0, /* tp_clear */ |
| 145 | 0, /* tp_richcompare */ |
| 146 | 0, /* tp_weaklistoffset */ |
Raymond Hettinger | 1da1dbf | 2003-03-17 19:46:11 +0000 | [diff] [blame] | 147 | PyObject_SelfIter, /* tp_iter */ |
Guido van Rossum | 7dab242 | 2002-04-26 19:40:56 +0000 | [diff] [blame] | 148 | (iternextfunc)enum_next, /* tp_iternext */ |
Guido van Rossum | ca5ed5b | 2002-07-16 21:02:42 +0000 | [diff] [blame] | 149 | 0, /* tp_methods */ |
Guido van Rossum | 7dab242 | 2002-04-26 19:40:56 +0000 | [diff] [blame] | 150 | 0, /* tp_members */ |
| 151 | 0, /* tp_getset */ |
| 152 | 0, /* tp_base */ |
| 153 | 0, /* tp_dict */ |
| 154 | 0, /* tp_descr_get */ |
| 155 | 0, /* tp_descr_set */ |
| 156 | 0, /* tp_dictoffset */ |
| 157 | 0, /* tp_init */ |
| 158 | PyType_GenericAlloc, /* tp_alloc */ |
| 159 | enum_new, /* tp_new */ |
| 160 | PyObject_GC_Del, /* tp_free */ |
| 161 | }; |