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

This commit is contained in:
Damien George 2014-03-26 18:37:06 +00:00
parent 38f0c607b0
commit 9e6e935df0
6 changed files with 58 additions and 29 deletions

View File

@ -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); 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); 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_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); void mp_obj_print(mp_obj_t o, mp_print_kind_t kind);

View File

@ -202,35 +202,38 @@ 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 // 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) { bool mp_obj_is_exception_type(mp_obj_t self_in) {
if (MP_OBJ_IS_TYPE(self_in, &mp_type_type)) { 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; mp_obj_type_t *self = self_in;
return self->make_new == mp_obj_exception_make_new; if (self->make_new == mp_obj_exception_make_new) {
} else { return true;
return false;
} }
} }
return mp_obj_is_subclass_fast(self_in, &mp_type_BaseException);
}
// return true if the given object is an instance of an exception type // 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) { 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) { void mp_obj_exception_clear_traceback(mp_obj_t self_in) {
// make sure self_in is an exception instance // make sure self_in is an exception instance
assert(mp_obj_get_type(self_in)->make_new == mp_obj_exception_make_new); // 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; mp_obj_exception_t *self = self_in;
// just set the traceback to the null object // just set the traceback to the null object
// we don't want to call any memory management functions here // we don't want to call any memory management functions here
self->traceback = MP_OBJ_NULL; self->traceback = MP_OBJ_NULL;
} }
}
void mp_obj_exception_add_traceback(mp_obj_t self_in, qstr file, machine_uint_t line, qstr block) { 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 // make sure self_in is an exception instance
assert(mp_obj_get_type(self_in)->make_new == mp_obj_exception_make_new); // 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; 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 // for traceback, we are just using the list object for convenience, it's not really a list of Python objects
@ -241,6 +244,7 @@ void mp_obj_exception_add_traceback(mp_obj_t self_in, qstr file, machine_uint_t
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)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)block);
} }
}
void mp_obj_exception_get_traceback(mp_obj_t self_in, machine_uint_t *n, machine_uint_t **values) { void mp_obj_exception_get_traceback(mp_obj_t self_in, machine_uint_t *n, machine_uint_t **values) {
// make sure self_in is an exception instance // make sure self_in is an exception instance

View File

@ -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 // object and classinfo should be type objects
// (but the function will fail gracefully if they are not) // (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 (;;) { for (;;) {
if (object == classinfo) { if (object == classinfo) {
return true; return true;
@ -482,7 +482,7 @@ bool mp_obj_is_subclass_fast(mp_obj_t object, mp_obj_t classinfo) {
return false; return false;
} }
mp_obj_type_t *self = object; const mp_obj_type_t *self = object;
// for a const struct, this entry might be NULL // for a const struct, this entry might be NULL
if (self->bases_tuple == MP_OBJ_NULL) { if (self->bases_tuple == MP_OBJ_NULL) {

View File

@ -17,6 +17,7 @@ Q(__repl_print__)
Q(__bool__) Q(__bool__)
Q(__len__) Q(__len__)
Q(__iter__)
Q(__getitem__) Q(__getitem__)
Q(__setitem__) Q(__setitem__)
Q(__add__) Q(__add__)

View File

@ -925,8 +925,13 @@ mp_obj_t rt_getiter(mp_obj_t o_in) {
if (type->getiter != NULL) { if (type->getiter != NULL) {
return type->getiter(o_in); return type->getiter(o_in);
} else { } else {
// check for __getitem__ method // check for __iter__ method
mp_obj_t dest[2]; mp_obj_t dest[2];
rt_load_method_maybe(o_in, MP_QSTR___iter__, dest);
if (dest[0] != MP_OBJ_NULL) {
// __iter__ exists, call it and return its result
return rt_call_method_n_kw(0, 0, dest);
} else {
rt_load_method_maybe(o_in, MP_QSTR___getitem__, dest); rt_load_method_maybe(o_in, MP_QSTR___getitem__, dest);
if (dest[0] != MP_OBJ_NULL) { if (dest[0] != MP_OBJ_NULL) {
// __getitem__ exists, create an iterator // __getitem__ exists, create an iterator
@ -937,15 +942,24 @@ mp_obj_t rt_getiter(mp_obj_t o_in) {
} }
} }
} }
}
mp_obj_t rt_iternext(mp_obj_t o_in) { mp_obj_t rt_iternext(mp_obj_t o_in) {
mp_obj_type_t *type = mp_obj_get_type(o_in); mp_obj_type_t *type = mp_obj_get_type(o_in);
if (type->iternext != NULL) { if (type->iternext != NULL) {
return type->iternext(o_in); return type->iternext(o_in);
} else {
// 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 { } 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))); nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "'%s' object is not an iterator", mp_obj_get_type_str(o_in)));
} }
} }
}
mp_obj_t rt_make_raise_obj(mp_obj_t o) { mp_obj_t rt_make_raise_obj(mp_obj_t o) {
DEBUG_printf("raise %p\n", o); DEBUG_printf("raise %p\n", o);

10
py/vm.c
View File

@ -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 // outer exception handling loop
for (;;) { for (;;) {
outer_dispatch_loop:
if (nlr_push(&nlr) == 0) { if (nlr_push(&nlr) == 0) {
// If we have exception to inject, now that we finish setting up // If we have exception to inject, now that we finish setting up
// execution context, raise it. This works as if RAISE_VARARGS // execution context, raise it. This works as if RAISE_VARARGS
@ -642,6 +643,15 @@ unwind_return:
} else { } else {
// exception occurred // 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 // set file and line number that the exception occurred at
// TODO: don't set traceback for exceptions re-raised by END_FINALLY. // TODO: don't set traceback for exceptions re-raised by END_FINALLY.
// But consider how to handle nested exceptions. // But consider how to handle nested exceptions.