summaryrefslogtreecommitdiffstatshomepage
diff options
context:
space:
mode:
-rw-r--r--py/obj.h2
-rw-r--r--py/objexcept.c14
-rw-r--r--py/objgenerator.c69
-rw-r--r--py/objgenerator.h1
-rw-r--r--tests/basics/generator-return.py10
-rw-r--r--tests/basics/generator_close.py59
6 files changed, 143 insertions, 12 deletions
diff --git a/py/obj.h b/py/obj.h
index f80bd43523..9599d137dd 100644
--- a/py/obj.h
+++ b/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);
diff --git a/py/objexcept.c b/py/objexcept.c
index de9bf1694f..f96523ed23 100644
--- a/py/objexcept.c
+++ b/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);
diff --git a/py/objgenerator.c b/py/objgenerator.c
index 39e362e041..3cfa02a0a0 100644
--- a/py/objgenerator.c
+++ b/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;
+
+ case MP_VM_RETURN_YIELD:
+ return *self->sp;
+
+ case MP_VM_RETURN_EXCEPTION:
+ self->ip = 0;
+ 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 {
- // TODO return StopIteration with value *self->sp
- return mp_const_stop_iteration;
+ nlr_jump(mp_obj_new_exception_args(&mp_type_StopIteration, 1, &ret));
}
case MP_VM_RETURN_YIELD:
- return *self->sp;
+ return ret;
case MP_VM_RETURN_EXCEPTION:
- self->ip = 0;
- nlr_jump(self->state[self->n_state - 1]);
+ 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
diff --git a/py/objgenerator.h b/py/objgenerator.h
new file mode 100644
index 0000000000..3dc69aa0ef
--- /dev/null
+++ b/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);
diff --git a/tests/basics/generator-return.py b/tests/basics/generator-return.py
new file mode 100644
index 0000000000..a3ac88575e
--- /dev/null
+++ b/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))
diff --git a/tests/basics/generator_close.py b/tests/basics/generator_close.py
new file mode 100644
index 0000000000..aa563f2a8a
--- /dev/null
+++ b/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")