Merge pull request #399 from pfalcon/gen-defargs

objgenerator: Handle default args to generator functions.

1 files changed, 36 insertions, 0 deletions

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);