Merge pull request adafruit#399 from pfalcon/gen-defargs

dpgeorge · dpgeorge · commit 0aa4379543b4 · 2014-03-30T21:51:49.000+01:00
objgenerator: Handle default args to generator functions.
diff --git 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);
@@ -463,6 +462,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
@@ -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
@@ -18,21 +18,21 @@ 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;
     assert(MP_OBJ_IS_TYPE(self_fun, &mp_type_fun_bc));
     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 = {
@@ -227,7 +227,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);
@@ -254,9 +254,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
@@ -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)))
