From 2e482cdb7b1e457053edcb8585414ecd1eba2cdc Mon Sep 17 00:00:00 2001 From: Damien George Date: Sun, 16 Feb 2014 00:01:29 +0000 Subject: [PATCH] py: Implement *vargs support. Addresses issue #295. --- py/obj.h | 2 +- py/objfun.c | 57 ++++++++++++++++++++++++++++++------- py/runtime.c | 2 +- tests/basics/fun-varargs.py | 41 ++++++++++++++++++++++++++ 4 files changed, 89 insertions(+), 13 deletions(-) create mode 100644 tests/basics/fun-varargs.py diff --git a/py/obj.h b/py/obj.h index dd6d217e4d..55be768730 100644 --- a/py/obj.h +++ b/py/obj.h @@ -236,7 +236,7 @@ mp_obj_t mp_obj_new_exception_msg(const mp_obj_type_t *exc_type, const char *msg mp_obj_t mp_obj_new_exception_msg_varg(const mp_obj_type_t *exc_type, const char *fmt, ...); // counts args by number of % symbols in fmt, excluding %%; can only handle void* sizes (ie no float/double!) mp_obj_t mp_obj_new_range(int start, int stop, int step); mp_obj_t mp_obj_new_range_iterator(int cur, int stop, int step); -mp_obj_t mp_obj_new_fun_bc(int n_args, mp_obj_t def_args, uint n_state, const byte *code); +mp_obj_t mp_obj_new_fun_bc(uint scope_flags, uint n_args, mp_obj_t def_args, uint n_state, 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, uint n_state, int n_args, const mp_obj_t *args); diff --git a/py/objfun.c b/py/objfun.c index 0c7ccf7984..73fd751a0d 100644 --- a/py/objfun.c +++ b/py/objfun.c @@ -10,6 +10,7 @@ #include "obj.h" #include "objtuple.h" #include "map.h" +#include "runtime0.h" #include "runtime.h" #include "bc.h" @@ -132,28 +133,52 @@ mp_obj_t rt_make_function_var_between(int n_args_min, int n_args_max, mp_fun_var typedef struct _mp_obj_fun_bc_t { mp_obj_base_t base; mp_map_t *globals; // the context within which this function was defined - short n_args; // number of arguments this function takes - short n_def_args; // number of default arguments + struct { + machine_uint_t n_args : 15; // number of arguments this function takes + machine_uint_t n_def_args : 15; // number of default arguments + machine_uint_t takes_var_args : 1; // set if this function takes variable args + machine_uint_t takes_kw_args : 1; // set if this function takes keyword args + }; uint n_state; // total state size for the executing function (incl args, locals, stack) const byte *bytecode; // bytecode for the function - mp_obj_t def_args[]; // values of default args, if any + mp_obj_t extra_args[]; // values of default args (if any), plus a slot at the end for var args (if it takes them) } mp_obj_fun_bc_t; STATIC mp_obj_t fun_bc_call(mp_obj_t self_in, uint n_args, uint n_kw, const mp_obj_t *args) { mp_obj_fun_bc_t *self = self_in; - if (n_args < self->n_args - self->n_def_args || n_args > self->n_args) { - 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)); + mp_obj_t *extra_args = self->extra_args + self->n_def_args; + uint n_extra_args = 0; + + if (n_args > self->n_args) { + // given more than enough arguments + if (!self->takes_var_args) { + goto arg_error; + } + // put extra arguments in varargs tuple + *extra_args = mp_obj_new_tuple(n_args - self->n_args, args + self->n_args); + n_extra_args = 1; + n_args = self->n_args; + } else if (n_args >= self->n_args - self->n_def_args) { + // given enough arguments, but may need to use some default arguments + if (self->takes_var_args) { + *extra_args = mp_const_empty_tuple; + n_extra_args = 1; + } + extra_args -= self->n_args - n_args; + n_extra_args += self->n_args - n_args; + } else { + goto arg_error; } + if (n_kw != 0) { nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "function does not take keyword arguments")); } - uint use_def_args = self->n_args - n_args; mp_map_t *old_globals = rt_globals_get(); rt_globals_set(self->globals); mp_obj_t result; - mp_vm_return_kind_t vm_return_kind = mp_execute_byte_code(self->bytecode, args, n_args, self->def_args + self->n_def_args - use_def_args, use_def_args, self->n_state, &result); + mp_vm_return_kind_t vm_return_kind = mp_execute_byte_code(self->bytecode, args, n_args, extra_args, n_extra_args, self->n_state, &result); rt_globals_set(old_globals); if (vm_return_kind == MP_VM_RETURN_NORMAL) { @@ -161,6 +186,9 @@ STATIC mp_obj_t fun_bc_call(mp_obj_t self_in, uint n_args, uint n_kw, const mp_o } else { // MP_VM_RETURN_EXCEPTION nlr_jump(result); } + +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)); } const mp_obj_type_t fun_bc_type = { @@ -169,21 +197,28 @@ const mp_obj_type_t fun_bc_type = { .call = fun_bc_call, }; -mp_obj_t mp_obj_new_fun_bc(int n_args, mp_obj_t def_args_in, uint n_state, const byte *code) { - int n_def_args = 0; +mp_obj_t mp_obj_new_fun_bc(uint scope_flags, uint n_args, mp_obj_t def_args_in, uint n_state, const byte *code) { + uint n_def_args = 0; + uint n_extra_args = 0; mp_obj_tuple_t *def_args = def_args_in; if (def_args != MP_OBJ_NULL) { n_def_args = def_args->len; + n_extra_args = def_args->len; } - mp_obj_fun_bc_t *o = m_new_obj_var(mp_obj_fun_bc_t, mp_obj_t, n_def_args); + if ((scope_flags & MP_SCOPE_FLAG_VARARGS) != 0) { + n_extra_args += 1; + } + mp_obj_fun_bc_t *o = m_new_obj_var(mp_obj_fun_bc_t, mp_obj_t, n_extra_args); o->base.type = &fun_bc_type; o->globals = rt_globals_get(); o->n_args = n_args; o->n_def_args = n_def_args; + o->takes_var_args = (scope_flags & MP_SCOPE_FLAG_VARARGS) != 0; + o->takes_kw_args = (scope_flags & MP_SCOPE_FLAG_VARKEYWORDS) != 0; o->n_state = n_state; o->bytecode = code; if (def_args != MP_OBJ_NULL) { - memcpy(o->def_args, def_args->items, n_def_args * sizeof(*o->def_args)); + memcpy(o->extra_args, def_args->items, n_def_args * sizeof(mp_obj_t)); } return o; } diff --git a/py/runtime.c b/py/runtime.c index 798f7b671c..9d93dbf2ad 100644 --- a/py/runtime.c +++ b/py/runtime.c @@ -714,7 +714,7 @@ mp_obj_t rt_make_function_from_id(int unique_code_id, mp_obj_t def_args) { mp_obj_t fun; switch (c->kind) { case MP_CODE_BYTE: - fun = mp_obj_new_fun_bc(c->n_args, def_args, c->n_state, c->u_byte.code); + fun = mp_obj_new_fun_bc(c->scope_flags, c->n_args, def_args, c->n_state, c->u_byte.code); break; case MP_CODE_NATIVE: fun = rt_make_function_n(c->n_args, c->u_native.fun); diff --git a/tests/basics/fun-varargs.py b/tests/basics/fun-varargs.py new file mode 100644 index 0000000000..ff29c33f7b --- /dev/null +++ b/tests/basics/fun-varargs.py @@ -0,0 +1,41 @@ +# function with just varargs +def f1(*args): + print(args) + +f1() +f1(1) +f1(1, 2) + +# function with 1 arg, then varargs +def f2(a, *args): + print(a, args) + +f2(1) +f2(1, 2) +f2(1, 2, 3) + +# function with 2 args, then varargs +def f3(a, b, *args): + print(a, b, args) + +f3(1, 2) +f3(1, 2, 3) +f3(1, 2, 3, 4) + +# function with 1 default arg, then varargs +def f4(a=0, *args): + print(a, args) + +f4() +f4(1) +f4(1, 2) +f4(1, 2, 3) + +# function with 1 arg, 1 default arg, then varargs +def f5(a, b=0, *args): + print(a, b, args) + +f5(1) +f5(1, 2) +f5(1, 2, 3) +f5(1, 2, 3, 4)