From 7fafb28f6d5b2a79f02a59b9b3bdea7ad920167e Mon Sep 17 00:00:00 2001 From: Paul Sokolovsky Date: Sun, 30 Mar 2014 20:21:28 +0300 Subject: [PATCH] objgenerator: Handle default args to generator functions. Addresses #397. --- py/obj.h | 3 ++- py/objfun.c | 36 ++++++++++++++++++++++++++++++++++ py/objgenerator.c | 21 +++++++++++--------- tests/basics/generator-args.py | 9 +++++++++ 4 files changed, 59 insertions(+), 10 deletions(-) create mode 100644 tests/basics/generator-args.py diff --git a/py/obj.h b/py/obj.h index 88d9253ec..20e2e5eee 100644 --- a/py/obj.h +++ b/py/obj.h @@ -331,7 +331,6 @@ mp_obj_t mp_obj_new_range_iterator(int cur, int stop, int step); mp_obj_t mp_obj_new_fun_bc(uint scope_flags, qstr *args, uint n_args, mp_obj_t def_args, const byte *code); mp_obj_t mp_obj_new_fun_asm(uint n_args, void *fun); mp_obj_t mp_obj_new_gen_wrap(mp_obj_t fun); -mp_obj_t mp_obj_new_gen_instance(const byte *bytecode, int n_args, const mp_obj_t *args); mp_obj_t mp_obj_new_closure(mp_obj_t fun, mp_obj_t closure_tuple); mp_obj_t mp_obj_new_tuple(uint n, const mp_obj_t *items); mp_obj_t mp_obj_new_list(uint n, mp_obj_t *items); @@ -462,6 +461,8 @@ typedef struct _mp_obj_fun_native_t { // need this so we can define const object } mp_obj_fun_native_t; void mp_obj_fun_bc_get(mp_obj_t self_in, int *n_args, const byte **code); +bool mp_obj_fun_prepare_simple_args(mp_obj_t self_in, uint n_args, uint n_kw, const mp_obj_t *args, + uint *out_args1_len, const mp_obj_t **out_args1, uint *out_args2_len, const mp_obj_t **out_args2); mp_obj_t mp_identity(mp_obj_t self); MP_DECLARE_CONST_FUN_OBJ(mp_identity_obj); diff --git a/py/objfun.c b/py/objfun.c index 019101bed..154630afb 100644 --- a/py/objfun.c +++ b/py/objfun.c @@ -161,6 +161,42 @@ void dump_args(const mp_obj_t *a, int sz) { #endif } +// If it's possible to call a function without allocating new argument array, +// this function returns true, together with pointers to 2 subarrays to be used +// as arguments. Otherwise, it returns false. It is expected that this fucntion +// will be accompanied by another, mp_obj_fun_prepare_full_args(), which will +// instead take pointer to full-length out-array, and will fill it in. Rationale +// being that a caller can try this function and if it succeeds, the function call +// can be made without allocating extra memory. Otherwise, caller can allocate memory +// and try "full" function. These functions are expected to be refactoring of +// code in fun_bc_call() and evenrually replace it. +bool mp_obj_fun_prepare_simple_args(mp_obj_t self_in, uint n_args, uint n_kw, const mp_obj_t *args, + uint *out_args1_len, const mp_obj_t **out_args1, uint *out_args2_len, const mp_obj_t **out_args2) { + mp_obj_fun_bc_t *self = self_in; + + assert(n_kw == 0); + assert(self->takes_var_args == 0); + assert(self->takes_kw_args == 0); + + mp_obj_t *extra_args = self->extra_args + self->n_def_args; + uint n_extra_args = 0; + + if (n_args > self->n_args) { + goto arg_error; + } else { + extra_args -= self->n_args - n_args; + n_extra_args += self->n_args - n_args; + } + *out_args1 = args; + *out_args1_len = n_args; + *out_args2 = extra_args; + *out_args2_len = n_extra_args; + return true; + +arg_error: + nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "function takes %d positional arguments but %d were given", self->n_args, n_args)); +} + STATIC mp_obj_t fun_bc_call(mp_obj_t self_in, uint n_args, uint n_kw, const mp_obj_t *args) { DEBUG_printf("Input: "); dump_args(args, n_args); diff --git a/py/objgenerator.c b/py/objgenerator.c index d1bae30de..0b342dc50 100644 --- a/py/objgenerator.c +++ b/py/objgenerator.c @@ -18,6 +18,8 @@ typedef struct _mp_obj_gen_wrap_t { mp_obj_t *fun; } mp_obj_gen_wrap_t; +mp_obj_t mp_obj_new_gen_instance(const byte *bytecode, uint n_args, const mp_obj_t *args, uint n_args2, const mp_obj_t *args2); + STATIC mp_obj_t gen_wrap_call(mp_obj_t self_in, uint n_args, uint n_kw, const mp_obj_t *args) { mp_obj_gen_wrap_t *self = self_in; mp_obj_t self_fun = self->fun; @@ -25,14 +27,12 @@ STATIC mp_obj_t gen_wrap_call(mp_obj_t self_in, uint n_args, uint n_kw, const mp int bc_n_args; const byte *bc_code; mp_obj_fun_bc_get(self_fun, &bc_n_args, &bc_code); - if (n_args != bc_n_args) { - nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "function takes %d positional arguments but %d were given", bc_n_args, n_args)); - } - if (n_kw != 0) { - nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "function does not take keyword arguments")); - } - return mp_obj_new_gen_instance(bc_code, n_args, args); + const mp_obj_t *args1, *args2; + uint len1, len2; + assert(mp_obj_fun_prepare_simple_args(self_fun, n_args, n_kw, args, &len1, &args1, &len2, &args2)); + + return mp_obj_new_gen_instance(bc_code, len1, args1, len2, args2); } const mp_obj_type_t mp_type_gen_wrap = { @@ -220,7 +220,7 @@ const mp_obj_type_t mp_type_gen_instance = { .locals_dict = (mp_obj_t)&gen_instance_locals_dict, }; -mp_obj_t mp_obj_new_gen_instance(const byte *bytecode, int n_args, const mp_obj_t *args) { +mp_obj_t mp_obj_new_gen_instance(const byte *bytecode, uint n_args, const mp_obj_t *args, uint n_args2, const mp_obj_t *args2) { const byte *code_info = bytecode; // get code info size, and skip the line number table machine_uint_t code_info_size = bytecode[0] | (bytecode[1] << 8) | (bytecode[2] << 16) | (bytecode[3] << 24); @@ -247,9 +247,12 @@ mp_obj_t mp_obj_new_gen_instance(const byte *bytecode, int n_args, const mp_obj_ o->n_state = n_state; // copy args to end of state array, in reverse (that's how mp_execute_byte_code_2 needs it) - for (int i = 0; i < n_args; i++) { + for (uint i = 0; i < n_args; i++) { o->state[n_state - 1 - i] = args[i]; } + for (uint i = 0; i < n_args2; i++) { + o->state[n_state - 1 - n_args - i] = args2[i]; + } return o; } diff --git a/tests/basics/generator-args.py b/tests/basics/generator-args.py new file mode 100644 index 000000000..70a1b5ada --- /dev/null +++ b/tests/basics/generator-args.py @@ -0,0 +1,9 @@ +# Handling of "complicated" arg forms to generators +# https://github.com/micropython/micropython/issues/397 +def gen(v=5): + for i in range(v): + yield i + +print(list(gen())) +# Still not supported, ditto for *args and **kwargs +#print(list(gen(v=10)))