From 07205ec323ca98d1c23af51e099119954d0fe7a4 Mon Sep 17 00:00:00 2001 From: "John R. Lenton" Date: Mon, 13 Jan 2014 02:31:00 +0000 Subject: [PATCH] added zip() --- py/misc.h | 1 + py/mpqstrraw.h | 1 + py/obj.h | 5 ++++ py/objtuple.c | 20 ++++++++++--- py/objzip.c | 59 +++++++++++++++++++++++++++++++++++++++ py/py.mk | 1 + py/runtime.c | 1 + tests/basics/tests/zip.py | 2 ++ 8 files changed, 86 insertions(+), 4 deletions(-) create mode 100644 py/objzip.c create mode 100644 tests/basics/tests/zip.py diff --git a/py/misc.h b/py/misc.h index 1bf4d8f291..22158c71a7 100644 --- a/py/misc.h +++ b/py/misc.h @@ -21,6 +21,7 @@ typedef unsigned int uint; #define m_renew(type, ptr, old_num, new_num) ((type*)(m_realloc((ptr), sizeof(type) * (old_num), sizeof(type) * (new_num)))) #define m_del(type, ptr, num) m_free(ptr, sizeof(type) * (num)) #define m_del_obj(type, ptr) (m_del(type, ptr, 1)) +#define m_del_var(obj_type, var_type, var_num, ptr) (m_free(ptr, sizeof(obj_type) + sizeof(var_type) * (var_num))) void *m_malloc(int num_bytes); void *m_malloc0(int num_bytes); diff --git a/py/mpqstrraw.h b/py/mpqstrraw.h index e4404a36f0..e33194fd45 100644 --- a/py/mpqstrraw.h +++ b/py/mpqstrraw.h @@ -60,6 +60,7 @@ Q(set) Q(sum) Q(tuple) Q(type) +Q(zip) Q(append) Q(pop) diff --git a/py/obj.h b/py/obj.h index a9a5827ee7..8a2cba62bb 100644 --- a/py/obj.h +++ b/py/obj.h @@ -285,6 +285,7 @@ mp_obj_t mp_obj_complex_binary_op(int op, mp_float_t lhs_real, mp_float_t lhs_im // tuple extern const mp_obj_type_t tuple_type; void mp_obj_tuple_get(mp_obj_t self_in, uint *len, mp_obj_t **items); +void mp_obj_tuple_del(mp_obj_t self_in); // list extern const mp_obj_type_t list_type; @@ -306,6 +307,10 @@ void mp_obj_set_store(mp_obj_t self_in, mp_obj_t item); extern const mp_obj_type_t slice_type; void mp_obj_slice_get(mp_obj_t self_in, machine_int_t *start, machine_int_t *stop, machine_int_t *step); +// zip +extern const mp_obj_type_t zip_type; + + // functions typedef struct _mp_obj_fun_native_t { // need this so we can define const objects (to go in ROM) mp_obj_base_t base; diff --git a/py/objtuple.c b/py/objtuple.c index 7685cc449f..5c1169ed55 100644 --- a/py/objtuple.c +++ b/py/objtuple.c @@ -116,8 +116,10 @@ mp_obj_t mp_obj_new_tuple(uint n, const mp_obj_t *items) { mp_obj_tuple_t *o = m_new_obj_var(mp_obj_tuple_t, mp_obj_t, n); o->base.type = &tuple_type; o->len = n; - for (int i = 0; i < n; i++) { - o->items[i] = items[i]; + if (items) { + for (int i = 0; i < n; i++) { + o->items[i] = items[i]; + } } return o; } @@ -138,8 +140,18 @@ mp_obj_t mp_obj_new_tuple_reverse(uint n, const mp_obj_t *items) { void mp_obj_tuple_get(mp_obj_t self_in, uint *len, mp_obj_t **items) { assert(MP_OBJ_IS_TYPE(self_in, &tuple_type)); mp_obj_tuple_t *self = self_in; - *len = self->len; - *items = &self->items[0]; + if (len) { + *len = self->len; + } + if (items) { + *items = &self->items[0]; + } +} + +void mp_obj_tuple_del(mp_obj_t self_in) { + assert(MP_OBJ_IS_TYPE(self_in, &tuple_type)); + mp_obj_tuple_t *self = self_in; + m_del_var(mp_obj_tuple_t, mp_obj_t, self->len, self); } /******************************************************************************/ diff --git a/py/objzip.c b/py/objzip.c new file mode 100644 index 0000000000..2ab80af0b5 --- /dev/null +++ b/py/objzip.c @@ -0,0 +1,59 @@ +#include +#include + +#include "misc.h" +#include "mpconfig.h" +#include "obj.h" +#include "runtime.h" + +typedef struct _mp_obj_zip_t { + mp_obj_base_t base; + int n_iters; + mp_obj_t iters[]; +} mp_obj_zip_t; + +static mp_obj_t zip_getiter(mp_obj_t self_in) { + return self_in; +} + +static mp_obj_t zip_iternext(mp_obj_t self_in); + +static mp_obj_t zip_make_new(mp_obj_t type_in, int n_args, const mp_obj_t *args) { + /* NOTE: args are backwards */ + mp_obj_zip_t *o = m_new_obj_var(mp_obj_zip_t, mp_obj_t, n_args); + o->base.type = &zip_type; + o->n_iters = n_args; + for (int i = 0; i < n_args; i++) { + o->iters[i] = rt_getiter(args[n_args-i-1]); + } + return o; +} + +const mp_obj_type_t zip_type = { + { &mp_const_type }, + "zip", + .make_new = zip_make_new, + .iternext = zip_iternext, + .getiter = zip_getiter, +}; + +static mp_obj_t zip_iternext(mp_obj_t self_in) { + assert(MP_OBJ_IS_TYPE(self_in, &zip_type)); + mp_obj_zip_t *self = self_in; + mp_obj_t *items; + if (self->n_iters == 0) { + return mp_const_stop_iteration; + } + mp_obj_t o = mp_obj_new_tuple(self->n_iters, NULL); + mp_obj_tuple_get(o, NULL, &items); + + for (int i = 0; i < self->n_iters; i++) { + mp_obj_t next = rt_iternext(self->iters[i]); + if (next == mp_const_stop_iteration) { + mp_obj_tuple_del(o); + return mp_const_stop_iteration; + } + items[i] = next; + } + return o; +} diff --git a/py/py.mk b/py/py.mk index a44a8bad02..95f9c07671 100644 --- a/py/py.mk +++ b/py/py.mk @@ -97,6 +97,7 @@ PY_O_BASENAME = \ vm.o \ showbc.o \ repl.o \ + objzip.o \ # prepend the build destination prefix to the py object files diff --git a/py/runtime.c b/py/runtime.c index 2af86b6abd..02f965dd63 100644 --- a/py/runtime.c +++ b/py/runtime.c @@ -113,6 +113,7 @@ void rt_init(void) { mp_map_add_qstr(&map_builtins, MP_QSTR_set, (mp_obj_t)&set_type); mp_map_add_qstr(&map_builtins, MP_QSTR_tuple, (mp_obj_t)&tuple_type); mp_map_add_qstr(&map_builtins, MP_QSTR_type, (mp_obj_t)&mp_const_type); + mp_map_add_qstr(&map_builtins, MP_QSTR_zip, (mp_obj_t)&zip_type); // built-in user functions; TODO covert all to &mp_builtin_xxx's mp_map_add_qstr(&map_builtins, MP_QSTR_abs, rt_make_function_1(mp_builtin_abs)); diff --git a/tests/basics/tests/zip.py b/tests/basics/tests/zip.py new file mode 100644 index 0000000000..c0109094f4 --- /dev/null +++ b/tests/basics/tests/zip.py @@ -0,0 +1,2 @@ +print(list(zip())) +print(list(zip([1], {2,3})))