objgenerator: Implement return with value and .close() method.

Return with value gets converted to StopIteration(value). Implementation
keeps optimizing against creating of possibly unneeded exception objects,
so there're considerable refactoring to implement these features.

py/obj.h

@@ -273,6 +273,7 @@ mp_obj_t mp_obj_new_float(mp_float_t val);
 mp_obj_t mp_obj_new_complex(mp_float_t real, mp_float_t imag);
 #endif
 mp_obj_t mp_obj_new_exception(const mp_obj_type_t *exc_type);
+mp_obj_t mp_obj_new_exception_args(const mp_obj_type_t *exc_type, uint n_args, const mp_obj_t *args);
 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);
@@ -342,6 +343,7 @@ machine_int_t mp_obj_int_get_checked(mp_obj_t self_in);
 // exception
 bool mp_obj_is_exception_type(mp_obj_t self_in);
 bool mp_obj_is_exception_instance(mp_obj_t self_in);
+bool mp_obj_exception_match(mp_obj_t exc, const mp_obj_type_t *exc_type);
 void mp_obj_exception_clear_traceback(mp_obj_t self_in);
 void mp_obj_exception_add_traceback(mp_obj_t self_in, qstr file, machine_uint_t line, qstr block);
 void mp_obj_exception_get_traceback(mp_obj_t self_in, machine_uint_t *n, machine_uint_t **values);

py/objexcept.c

@@ -8,6 +8,8 @@
 #include "qstr.h"
 #include "obj.h"
 #include "objtuple.h"
+#include "runtime.h"
+#include "runtime0.h"
 
 // This is unified class for C-level and Python-level exceptions
 // Python-level exceptions have empty ->msg and all arguments are in
@@ -156,6 +158,11 @@ mp_obj_t mp_obj_new_exception(const mp_obj_type_t *exc_type) {
     return mp_obj_new_exception_msg_varg(exc_type, NULL);
 }
 
+mp_obj_t mp_obj_new_exception_args(const mp_obj_type_t *exc_type, uint n_args, const mp_obj_t *args) {
+    assert(exc_type->make_new == mp_obj_exception_make_new);
+    return exc_type->make_new((mp_obj_t)exc_type, n_args, 0, args);
+}
+
 mp_obj_t mp_obj_new_exception_msg(const mp_obj_type_t *exc_type, const char *msg) {
     return mp_obj_new_exception_msg_varg(exc_type, msg);
 }
@@ -202,6 +209,13 @@ bool mp_obj_is_exception_instance(mp_obj_t self_in) {
     return mp_obj_get_type(self_in)->make_new == mp_obj_exception_make_new;
 }
 
+// return true if exception (type or instance) is a subclass of given
+// exception type.
+bool mp_obj_exception_match(mp_obj_t exc, const mp_obj_type_t *exc_type) {
+    // TODO: move implementation from RT_BINARY_OP_EXCEPTION_MATCH here.
+    return rt_binary_op(RT_BINARY_OP_EXCEPTION_MATCH, exc, (mp_obj_t)exc_type) == mp_const_true;
+}
+
 void mp_obj_exception_clear_traceback(mp_obj_t self_in) {
     // make sure self_in is an exception instance
     assert(mp_obj_get_type(self_in)->make_new == mp_obj_exception_make_new);

py/objgenerator.c

@@ -8,6 +8,7 @@
 #include "obj.h"
 #include "runtime.h"
 #include "bc.h"
+#include "objgenerator.h"
 
 /******************************************************************************/
 /* generator wrapper                                                          */
@@ -73,9 +74,10 @@ mp_obj_t gen_instance_getiter(mp_obj_t self_in) {
     return self_in;
 }
 
-STATIC mp_obj_t gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value) {
+mp_obj_t mp_obj_gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value, mp_vm_return_kind_t *ret_kind) {
     mp_obj_gen_instance_t *self = self_in;
     if (self->ip == 0) {
+        *ret_kind = MP_VM_RETURN_NORMAL;
         return mp_const_stop_iteration;
     }
     if (self->sp == self->state - 1) {
@@ -85,10 +87,11 @@ STATIC mp_obj_t gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw
     } else {
         *self->sp = send_value;
     }
-    mp_vm_return_kind_t vm_return_kind = mp_execute_byte_code_2(self->code_info, &self->ip,
+    *ret_kind = mp_execute_byte_code_2(self->code_info, &self->ip,
         &self->state[self->n_state - 1], &self->sp, (mp_exc_stack*)(self->state + self->n_state),
         &self->exc_sp, throw_value);
-    switch (vm_return_kind) {
+
+    switch (*ret_kind) {
         case MP_VM_RETURN_NORMAL:
             // Explicitly mark generator as completed. If we don't do this,
             // subsequent next() may re-execute statements after last yield
@@ -96,19 +99,39 @@ STATIC mp_obj_t gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw
             // TODO: check how return with value behaves under such conditions
             // in CPython.
             self->ip = 0;
-            if (*self->sp == mp_const_none) {
+            return *self->sp;
-                return mp_const_stop_iteration;
-            } else {
-                // TODO return StopIteration with value *self->sp
-                return mp_const_stop_iteration;
-            }
 
         case MP_VM_RETURN_YIELD:
             return *self->sp;
 
         case MP_VM_RETURN_EXCEPTION:
             self->ip = 0;
-            nlr_jump(self->state[self->n_state - 1]);
+            return self->state[self->n_state - 1];
+
+        default:
+            assert(0);
+            return mp_const_none;
+    }
+}
+
+STATIC mp_obj_t gen_resume_and_raise(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value) {
+    mp_vm_return_kind_t ret_kind;
+    mp_obj_t ret = mp_obj_gen_resume(self_in, send_value, throw_value, &ret_kind);
+
+    switch (ret_kind) {
+        case MP_VM_RETURN_NORMAL:
+            // Optimize return w/o value in case generator is used in for loop
+            if (ret == mp_const_none) {
+                return mp_const_stop_iteration;
+            } else {
+                nlr_jump(mp_obj_new_exception_args(&mp_type_StopIteration, 1, &ret));
+            }
+
+        case MP_VM_RETURN_YIELD:
+            return ret;
+
+        case MP_VM_RETURN_EXCEPTION:
+            nlr_jump(ret);
 
         default:
             assert(0);
@@ -117,11 +140,11 @@ STATIC mp_obj_t gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw
 }
 
 mp_obj_t gen_instance_iternext(mp_obj_t self_in) {
-    return gen_resume(self_in, mp_const_none, MP_OBJ_NULL);
+    return gen_resume_and_raise(self_in, mp_const_none, MP_OBJ_NULL);
 }
 
 STATIC mp_obj_t gen_instance_send(mp_obj_t self_in, mp_obj_t send_value) {
-    mp_obj_t ret = gen_resume(self_in, send_value, MP_OBJ_NULL);
+    mp_obj_t ret = gen_resume_and_raise(self_in, send_value, MP_OBJ_NULL);
     if (ret == mp_const_stop_iteration) {
         nlr_jump(mp_obj_new_exception(&mp_type_StopIteration));
     } else {
@@ -132,7 +155,7 @@ STATIC mp_obj_t gen_instance_send(mp_obj_t self_in, mp_obj_t send_value) {
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(gen_instance_send_obj, gen_instance_send);
 
 STATIC mp_obj_t gen_instance_throw(uint n_args, const mp_obj_t *args) {
-    mp_obj_t ret = gen_resume(args[0], mp_const_none, n_args == 2 ? args[1] : args[2]);
+    mp_obj_t ret = gen_resume_and_raise(args[0], mp_const_none, n_args == 2 ? args[1] : args[2]);
     if (ret == mp_const_stop_iteration) {
         nlr_jump(mp_obj_new_exception(&mp_type_StopIteration));
     } else {
@@ -142,8 +165,30 @@ STATIC mp_obj_t gen_instance_throw(uint n_args, const mp_obj_t *args) {
 
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(gen_instance_throw_obj, 2, 4, gen_instance_throw);
 
+STATIC mp_obj_t gen_instance_close(mp_obj_t self_in) {
+    mp_vm_return_kind_t ret_kind;
+    mp_obj_t ret = mp_obj_gen_resume(self_in, mp_const_none, (mp_obj_t)&mp_type_GeneratorExit, &ret_kind);
+
+    if (ret_kind == MP_VM_RETURN_YIELD) {
+        nlr_jump(mp_obj_new_exception_msg(&mp_type_RuntimeError, "generator ignored GeneratorExit"));
+    }
+    // Swallow StopIteration & GeneratorExit (== successful close), and re-raise any other
+    if (ret_kind == MP_VM_RETURN_EXCEPTION) {
+        if (mp_obj_exception_match(ret, &mp_type_GeneratorExit) ||
+            mp_obj_exception_match(ret, &mp_type_StopIteration)) {
+            return mp_const_none;
+        }
+        nlr_jump(ret);
+    }
+
+    // The only choice left is MP_VM_RETURN_NORMAL which is successful close
+    return mp_const_none;
+}
+
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(gen_instance_close_obj, gen_instance_close);
 
 STATIC const mp_method_t gen_type_methods[] = {
+    { "close", &gen_instance_close_obj },
     { "send", &gen_instance_send_obj },
     { "throw", &gen_instance_throw_obj },
     { NULL, NULL }, // end-of-list sentinel

py/objgenerator.h

@@ -0,0 +1 @@
+mp_obj_t mp_obj_gen_resume(mp_obj_t self_in, mp_obj_t send_val, mp_obj_t throw_val, mp_vm_return_kind_t *ret_kind);

tests/basics/generator-return.py

@@ -0,0 +1,10 @@
+def gen():
+    yield 1
+    return 42
+
+g = gen()
+print(next(g))
+try:
+    print(next(g))
+except StopIteration as e:
+    print(repr(e))

tests/basics/generator_close.py

@@ -0,0 +1,59 @@
+def gen1():
+    yield 1
+    yield 2
+
+# Test that it's possible to close just created gen
+g = gen1()
+print(g.close())
+try:
+    next(g)
+except StopIteration:
+    print("StopIteration")
+
+# Test that it's possible to close gen in progress
+g = gen1()
+print(next(g))
+print(g.close())
+try:
+    next(g)
+    print("No StopIteration")
+except StopIteration:
+    print("StopIteration")
+
+# Test that it's possible to close finished gen
+g = gen1()
+print(list(g))
+print(g.close())
+try:
+    next(g)
+    print("No StopIteration")
+except StopIteration:
+    print("StopIteration")
+
+
+# Throwing StopIteration in response to close() is ok
+def gen2():
+    try:
+        yield 1
+        yield 2
+    except:
+        raise StopIteration
+
+g = gen2()
+next(g)
+print(g.close())
+
+# Any other exception is propagated to the .close() caller
+def gen3():
+    try:
+        yield 1
+        yield 2
+    except:
+        raise ValueError
+
+g = gen3()
+next(g)
+try:
+    print(g.close())
+except ValueError:
+    print("ValueError")