Commit 9e6e935d authored by Damien George's avatar Damien George
Browse files

py: Add support for user-defined iterators via __iter__, __next__.

parent 38f0c607
......@@ -261,7 +261,7 @@ mp_obj_t mp_obj_new_module(qstr module_name);
mp_obj_type_t *mp_obj_get_type(mp_obj_t o_in);
const char *mp_obj_get_type_str(mp_obj_t o_in);
bool mp_obj_is_subclass_fast(mp_obj_t object, mp_obj_t classinfo); // arguments should be type objects
bool mp_obj_is_subclass_fast(mp_const_obj_t object, mp_const_obj_t classinfo); // arguments should be type objects
void mp_obj_print_helper(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t o_in, mp_print_kind_t kind);
void mp_obj_print(mp_obj_t o, mp_print_kind_t kind);
......
......@@ -202,44 +202,48 @@ mp_obj_t mp_obj_new_exception_msg_varg(const mp_obj_type_t *exc_type, const char
}
// return true if the given object is an exception type
// TODO make this work for user defined exceptions
bool mp_obj_is_exception_type(mp_obj_t self_in) {
if (MP_OBJ_IS_TYPE(self_in, &mp_type_type)) {
// optimisation when self_in is a builtin exception
mp_obj_type_t *self = self_in;
return self->make_new == mp_obj_exception_make_new;
} else {
return false;
if (self->make_new == mp_obj_exception_make_new) {
return true;
}
}
return mp_obj_is_subclass_fast(self_in, &mp_type_BaseException);
}
// return true if the given object is an instance of an exception type
// TODO make this work for user defined exceptions
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 mp_obj_is_exception_type(mp_obj_get_type(self_in));
}
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);
mp_obj_exception_t *self = self_in;
// TODO add traceback information to user-defined exceptions (need proper builtin subclassing for that)
if (mp_obj_get_type(self_in)->make_new == mp_obj_exception_make_new) {
mp_obj_exception_t *self = self_in;
// just set the traceback to the null object
// we don't want to call any memory management functions here
self->traceback = MP_OBJ_NULL;
// just set the traceback to the null object
// we don't want to call any memory management functions here
self->traceback = MP_OBJ_NULL;
}
}
void mp_obj_exception_add_traceback(mp_obj_t self_in, qstr file, machine_uint_t line, qstr block) {
// make sure self_in is an exception instance
assert(mp_obj_get_type(self_in)->make_new == mp_obj_exception_make_new);
mp_obj_exception_t *self = self_in;
// TODO add traceback information to user-defined exceptions (need proper builtin subclassing for that)
if (mp_obj_get_type(self_in)->make_new == mp_obj_exception_make_new) {
mp_obj_exception_t *self = self_in;
// for traceback, we are just using the list object for convenience, it's not really a list of Python objects
if (self->traceback == MP_OBJ_NULL) {
self->traceback = mp_obj_new_list(0, NULL);
// for traceback, we are just using the list object for convenience, it's not really a list of Python objects
if (self->traceback == MP_OBJ_NULL) {
self->traceback = mp_obj_new_list(0, NULL);
}
mp_obj_list_append(self->traceback, (mp_obj_t)(machine_uint_t)file);
mp_obj_list_append(self->traceback, (mp_obj_t)(machine_uint_t)line);
mp_obj_list_append(self->traceback, (mp_obj_t)(machine_uint_t)block);
}
mp_obj_list_append(self->traceback, (mp_obj_t)(machine_uint_t)file);
mp_obj_list_append(self->traceback, (mp_obj_t)(machine_uint_t)line);
mp_obj_list_append(self->traceback, (mp_obj_t)(machine_uint_t)block);
}
void mp_obj_exception_get_traceback(mp_obj_t self_in, machine_uint_t *n, machine_uint_t **values) {
......
......@@ -469,7 +469,7 @@ mp_obj_t mp_obj_new_super(mp_obj_t type, mp_obj_t obj) {
// object and classinfo should be type objects
// (but the function will fail gracefully if they are not)
bool mp_obj_is_subclass_fast(mp_obj_t object, mp_obj_t classinfo) {
bool mp_obj_is_subclass_fast(mp_const_obj_t object, mp_const_obj_t classinfo) {
for (;;) {
if (object == classinfo) {
return true;
......@@ -482,7 +482,7 @@ bool mp_obj_is_subclass_fast(mp_obj_t object, mp_obj_t classinfo) {
return false;
}
mp_obj_type_t *self = object;
const mp_obj_type_t *self = object;
// for a const struct, this entry might be NULL
if (self->bases_tuple == MP_OBJ_NULL) {
......
......@@ -17,6 +17,7 @@ Q(__repl_print__)
Q(__bool__)
Q(__len__)
Q(__iter__)
Q(__getitem__)
Q(__setitem__)
Q(__add__)
......
......@@ -925,15 +925,21 @@ mp_obj_t rt_getiter(mp_obj_t o_in) {
if (type->getiter != NULL) {
return type->getiter(o_in);
} else {
// check for __getitem__ method
// check for __iter__ method
mp_obj_t dest[2];
rt_load_method_maybe(o_in, MP_QSTR___getitem__, dest);
rt_load_method_maybe(o_in, MP_QSTR___iter__, dest);
if (dest[0] != MP_OBJ_NULL) {
// __getitem__ exists, create an iterator
return mp_obj_new_getitem_iter(dest);
// __iter__ exists, call it and return its result
return rt_call_method_n_kw(0, 0, dest);
} else {
// object not iterable
nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "'%s' object is not iterable", mp_obj_get_type_str(o_in)));
rt_load_method_maybe(o_in, MP_QSTR___getitem__, dest);
if (dest[0] != MP_OBJ_NULL) {
// __getitem__ exists, create an iterator
return mp_obj_new_getitem_iter(dest);
} else {
// object not iterable
nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "'%s' object is not iterable", mp_obj_get_type_str(o_in)));
}
}
}
}
......@@ -943,7 +949,15 @@ mp_obj_t rt_iternext(mp_obj_t o_in) {
if (type->iternext != NULL) {
return type->iternext(o_in);
} else {
nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "'%s' object is not an iterator", mp_obj_get_type_str(o_in)));
// check for __next__ method
mp_obj_t dest[2];
rt_load_method_maybe(o_in, MP_QSTR___next__, dest);
if (dest[0] != MP_OBJ_NULL) {
// __next__ exists, call it and return its result
return rt_call_method_n_kw(0, 0, dest);
} else {
nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "'%s' object is not an iterator", mp_obj_get_type_str(o_in)));
}
}
}
......
......@@ -124,6 +124,7 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
// outer exception handling loop
for (;;) {
outer_dispatch_loop:
if (nlr_push(&nlr) == 0) {
// If we have exception to inject, now that we finish setting up
// execution context, raise it. This works as if RAISE_VARARGS
......@@ -642,6 +643,15 @@ unwind_return:
} else {
// exception occurred
// check if it's a StopIteration within a for block
if (*save_ip == MP_BC_FOR_ITER && mp_obj_is_subclass_fast(mp_obj_get_type(nlr.ret_val), &mp_type_StopIteration)) {
ip = save_ip + 1;
DECODE_ULABEL; // the jump offset if iteration finishes; for labels are always forward
--sp; // pop the exhausted iterator
ip += unum; // jump to after for-block
goto outer_dispatch_loop; // continue with dispatch loop
}
// set file and line number that the exception occurred at
// TODO: don't set traceback for exceptions re-raised by END_FINALLY.
// But consider how to handle nested exceptions.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment