blob: 564cb1c47420423e8ef594c6c174fa1a07bc8cd6 [file] [log] [blame]
John R. Lenton9daa7892014-01-14 23:55:01 +00001#include <stdlib.h>
2#include <assert.h>
3
4#include "misc.h"
5#include "mpconfig.h"
Damien George55baff42014-01-21 21:40:13 +00006#include "qstr.h"
John R. Lenton9daa7892014-01-14 23:55:01 +00007#include "obj.h"
8#include "runtime.h"
9
10typedef struct _mp_obj_enumerate_t {
11 mp_obj_base_t base;
12 mp_obj_t iter;
13 machine_int_t cur;
14} mp_obj_enumerate_t;
15
16static mp_obj_t enumerate_getiter(mp_obj_t self_in) {
17 return self_in;
18}
19
20static mp_obj_t enumerate_iternext(mp_obj_t self_in);
21
22/* TODO: enumerate is one of the ones that can take args or kwargs.
23 Sticking to args for now */
Damien George20006db2014-01-18 14:10:48 +000024static mp_obj_t enumerate_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_obj_t *args) {
John R. Lenton9daa7892014-01-14 23:55:01 +000025 assert(n_args > 0);
John R. Lenton9daa7892014-01-14 23:55:01 +000026 mp_obj_enumerate_t *o = m_new_obj(mp_obj_enumerate_t);
27 o->base.type = &enumerate_type;
28 o->iter = rt_getiter(args[0]);
Damien George20006db2014-01-18 14:10:48 +000029 o->cur = n_args > 1 ? mp_obj_get_int(args[1]) : 0;
John R. Lenton9daa7892014-01-14 23:55:01 +000030 return o;
31}
32
33const mp_obj_type_t enumerate_type = {
34 { &mp_const_type },
35 "enumerate",
36 .make_new = enumerate_make_new,
37 .iternext = enumerate_iternext,
38 .getiter = enumerate_getiter,
39};
40
41static mp_obj_t enumerate_iternext(mp_obj_t self_in) {
42 assert(MP_OBJ_IS_TYPE(self_in, &enumerate_type));
43 mp_obj_enumerate_t *self = self_in;
44 mp_obj_t next = rt_iternext(self->iter);
45 if (next == mp_const_stop_iteration) {
46 return mp_const_stop_iteration;
47 } else {
48 mp_obj_t items[] = {MP_OBJ_NEW_SMALL_INT(self->cur++), next};
49 return mp_obj_new_tuple(2, items);
50 }
51}