John R. Lenton | 9daa789 | 2014-01-14 23:55:01 +0000 | [diff] [blame] | 1 | #include <stdlib.h> |
| 2 | #include <assert.h> |
| 3 | |
| 4 | #include "misc.h" |
| 5 | #include "mpconfig.h" |
Damien George | 55baff4 | 2014-01-21 21:40:13 +0000 | [diff] [blame] | 6 | #include "qstr.h" |
John R. Lenton | 9daa789 | 2014-01-14 23:55:01 +0000 | [diff] [blame] | 7 | #include "obj.h" |
| 8 | #include "runtime.h" |
| 9 | |
| 10 | typedef 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 | |
| 16 | static mp_obj_t enumerate_getiter(mp_obj_t self_in) { |
| 17 | return self_in; |
| 18 | } |
| 19 | |
| 20 | static 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 George | 20006db | 2014-01-18 14:10:48 +0000 | [diff] [blame] | 24 | static mp_obj_t enumerate_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_obj_t *args) { |
John R. Lenton | 9daa789 | 2014-01-14 23:55:01 +0000 | [diff] [blame] | 25 | assert(n_args > 0); |
John R. Lenton | 9daa789 | 2014-01-14 23:55:01 +0000 | [diff] [blame] | 26 | 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 George | 20006db | 2014-01-18 14:10:48 +0000 | [diff] [blame] | 29 | o->cur = n_args > 1 ? mp_obj_get_int(args[1]) : 0; |
John R. Lenton | 9daa789 | 2014-01-14 23:55:01 +0000 | [diff] [blame] | 30 | return o; |
| 31 | } |
| 32 | |
| 33 | const 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 | |
| 41 | static 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 | } |