diff --git a/py/obj.h b/py/obj.h index 68d9588b4d..e122f5a2bf 100644 --- a/py/obj.h +++ b/py/obj.h @@ -39,7 +39,7 @@ typedef struct _mp_obj_base_t mp_obj_base_t; #define MP_OBJ_IS_SMALL_INT(o) ((((mp_small_int_t)(o)) & 1) != 0) #define MP_OBJ_IS_QSTR(o) ((((mp_small_int_t)(o)) & 3) == 2) #define MP_OBJ_IS_OBJ(o) ((((mp_small_int_t)(o)) & 3) == 0) -#define MP_OBJ_IS_TYPE(o, t) (MP_OBJ_IS_OBJ(o) && (((mp_obj_base_t*)(o))->type == (t))) +#define MP_OBJ_IS_TYPE(o, t) (MP_OBJ_IS_OBJ(o) && (((mp_obj_base_t*)(o))->type == (t))) // this does not work for checking a string, use below macro for that #define MP_OBJ_IS_STR(o) (MP_OBJ_IS_QSTR(o) || MP_OBJ_IS_TYPE(o, &str_type)) #define MP_OBJ_SMALL_INT_VALUE(o) (((mp_small_int_t)(o)) >> 1) @@ -231,6 +231,7 @@ mp_obj_t mp_obj_new_dict(int n_args); mp_obj_t mp_obj_new_set(int n_args, mp_obj_t *items); mp_obj_t mp_obj_new_slice(mp_obj_t start, mp_obj_t stop, mp_obj_t step); mp_obj_t mp_obj_new_bound_meth(mp_obj_t meth, mp_obj_t self); +mp_obj_t mp_obj_new_getitem_iter(mp_obj_t *args); mp_obj_t mp_obj_new_module(qstr module_name); mp_obj_type_t *mp_obj_get_type(mp_obj_t o_in); diff --git a/py/objgetitemiter.c b/py/objgetitemiter.c new file mode 100644 index 0000000000..40ed1a1520 --- /dev/null +++ b/py/objgetitemiter.c @@ -0,0 +1,53 @@ +#include +#include + +#include "nlr.h" +#include "misc.h" +#include "mpconfig.h" +#include "qstr.h" +#include "obj.h" +#include "runtime.h" + +// this is a wrapper object that is turns something that has a __getitem__ method into an iterator + +typedef struct _mp_obj_getitem_iter_t { + mp_obj_base_t base; + mp_obj_t args[3]; +} mp_obj_getitem_iter_t; + +static mp_obj_t it_iternext(mp_obj_t self_in) { + mp_obj_getitem_iter_t *self = self_in; + nlr_buf_t nlr; + if (nlr_push(&nlr) == 0) { + // try to get next item + mp_obj_t value = rt_call_method_n_kw(1, 0, self->args); + self->args[2] = MP_OBJ_NEW_SMALL_INT(MP_OBJ_SMALL_INT_VALUE(self->args[2]) + 1); + nlr_pop(); + return value; + } else { + // an exception was raised + if (MP_OBJ_IS_TYPE(nlr.ret_val, &exception_type) && mp_obj_exception_get_type(nlr.ret_val) == MP_QSTR_StopIteration) { + // return mp_const_stop_iteration instead of raising StopIteration + return mp_const_stop_iteration; + } else { + // re-raise exception + nlr_jump(nlr.ret_val); + } + } +} + +static const mp_obj_type_t it_type = { + { &mp_const_type }, + "iterator", + .iternext = it_iternext +}; + +// args are those returned from rt_load_method_maybe (ie either an attribute or a method) +mp_obj_t mp_obj_new_getitem_iter(mp_obj_t *args) { + mp_obj_getitem_iter_t *o = m_new_obj(mp_obj_getitem_iter_t); + o->base.type = &it_type; + o->args[0] = args[0]; + o->args[1] = args[1]; + o->args[2] = MP_OBJ_NEW_SMALL_INT(0); + return o; +} diff --git a/py/objstr.c b/py/objstr.c index 723eebd614..3a4d69cfcc 100644 --- a/py/objstr.c +++ b/py/objstr.c @@ -77,7 +77,7 @@ mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) { if (MP_OBJ_IS_SMALL_INT(rhs_in)) { uint index = mp_get_index(mp_obj_get_type(lhs_in), lhs_len, rhs_in); if (MP_OBJ_IS_TYPE(lhs_in, &bytes_type)) { - return MP_OBJ_NEW_SMALL_INT(lhs_data[index]); + return MP_OBJ_NEW_SMALL_INT((mp_small_int_t)lhs_data[index]); } else { return mp_obj_new_str(lhs_data + index, 1, true); } @@ -549,7 +549,7 @@ mp_obj_t bytes_it_iternext(mp_obj_t self_in) { mp_obj_str_it_t *self = self_in; GET_STR_DATA_LEN(self->str, str, len); if (self->cur < len) { - mp_obj_t o_out = MP_OBJ_NEW_SMALL_INT(str[self->cur]); + mp_obj_t o_out = MP_OBJ_NEW_SMALL_INT((mp_small_int_t)str[self->cur]); self->cur += 1; return o_out; } else { diff --git a/py/objtype.c b/py/objtype.c index 75755f4fb9..afa34aa0c3 100644 --- a/py/objtype.c +++ b/py/objtype.c @@ -117,8 +117,8 @@ static mp_obj_t class_make_new(mp_obj_t self_in, uint n_args, uint n_kw, const m } // TODO somehow replace const char * with a qstr -static const char *binary_op_method_name[] = { - [RT_BINARY_OP_SUBSCR] = "__getitem__", +static const qstr binary_op_method_name[] = { + [RT_BINARY_OP_SUBSCR] = MP_QSTR___getitem__, /* RT_BINARY_OP_OR, RT_BINARY_OP_XOR, @@ -126,8 +126,8 @@ static const char *binary_op_method_name[] = { RT_BINARY_OP_LSHIFT, RT_BINARY_OP_RSHIFT, */ - [RT_BINARY_OP_ADD] = "__add__", - [RT_BINARY_OP_SUBTRACT] = "__sub__", + [RT_BINARY_OP_ADD] = MP_QSTR___add__, + [RT_BINARY_OP_SUBTRACT] = MP_QSTR___sub__, /* RT_BINARY_OP_MULTIPLY, RT_BINARY_OP_FLOOR_DIVIDE, @@ -157,16 +157,16 @@ static const char *binary_op_method_name[] = { RT_COMPARE_OP_IS, RT_COMPARE_OP_IS_NOT, */ - [RT_COMPARE_OP_EXCEPTION_MATCH] = "__not_implemented__", + [RT_COMPARE_OP_EXCEPTION_MATCH] = MP_QSTR_, // not implemented, used to make sure array has full size }; static mp_obj_t class_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) { mp_obj_class_t *lhs = lhs_in; - const char *op_name = binary_op_method_name[op]; - if (op_name == NULL) { + qstr op_name = binary_op_method_name[op]; + if (op_name == 0) { return MP_OBJ_NULL; } - mp_obj_t member = mp_obj_class_lookup(lhs->base.type, QSTR_FROM_STR_STATIC(op_name)); + mp_obj_t member = mp_obj_class_lookup(lhs->base.type, op_name); if (member != MP_OBJ_NULL) { return rt_call_function_2(member, lhs_in, rhs_in); } else { diff --git a/py/py.mk b/py/py.mk index d4946e2838..742e6e398e 100644 --- a/py/py.mk +++ b/py/py.mk @@ -47,6 +47,7 @@ PY_O_BASENAME = \ objfloat.o \ objfun.o \ objgenerator.o \ + objgetitemiter.o \ objint.o \ objint_longlong.o \ objlist.o \ diff --git a/py/qstrdefs.h b/py/qstrdefs.h index f2c4dfd97f..e76efaf0e0 100644 --- a/py/qstrdefs.h +++ b/py/qstrdefs.h @@ -13,6 +13,10 @@ Q(__next__) Q(__qualname__) Q(__repl_print__) +Q(__getitem__) +Q(__add__) +Q(__sub__) + Q(micropython) Q(byte_code) Q(native) diff --git a/py/runtime.c b/py/runtime.c index 3d56cc87ba..976d23a6b9 100644 --- a/py/runtime.c +++ b/py/runtime.c @@ -87,7 +87,7 @@ void rt_init(void) { // init loaded modules table mp_map_init(&map_loaded_modules, 3); - // built-in exceptions (TODO, make these proper classes) + // built-in exceptions (TODO, make these proper classes, and const if possible) mp_map_add_qstr(&map_builtins, MP_QSTR_AttributeError, mp_obj_new_exception(MP_QSTR_AttributeError)); mp_map_add_qstr(&map_builtins, MP_QSTR_IndexError, mp_obj_new_exception(MP_QSTR_IndexError)); mp_map_add_qstr(&map_builtins, MP_QSTR_KeyError, mp_obj_new_exception(MP_QSTR_KeyError)); @@ -100,6 +100,7 @@ void rt_init(void) { mp_map_add_qstr(&map_builtins, MP_QSTR_OverflowError, mp_obj_new_exception(MP_QSTR_OverflowError)); mp_map_add_qstr(&map_builtins, MP_QSTR_OSError, mp_obj_new_exception(MP_QSTR_OSError)); mp_map_add_qstr(&map_builtins, MP_QSTR_AssertionError, mp_obj_new_exception(MP_QSTR_AssertionError)); + mp_map_add_qstr(&map_builtins, MP_QSTR_StopIteration, mp_obj_new_exception(MP_QSTR_StopIteration)); // built-in objects mp_map_add_qstr(&map_builtins, MP_QSTR_Ellipsis, mp_const_ellipsis); @@ -805,7 +806,7 @@ mp_obj_t rt_load_attr(mp_obj_t base, qstr attr) { // use load_method mp_obj_t dest[2]; rt_load_method(base, attr, dest); - if (dest[1] == NULL) { + if (dest[1] == MP_OBJ_NULL) { // load_method returned just a normal attribute return dest[0]; } else { @@ -814,9 +815,10 @@ mp_obj_t rt_load_attr(mp_obj_t base, qstr attr) { } } -void rt_load_method(mp_obj_t base, qstr attr, mp_obj_t *dest) { - DEBUG_OP_printf("load method %p.%s\n", base, qstr_str(attr)); - +// no attribute found, returns: dest[0] == MP_OBJ_NULL, dest[1] == MP_OBJ_NULL +// normal attribute found, returns: dest[0] == , dest[1] == MP_OBJ_NULL +// method attribute found, returns: dest[0] == , dest[1] == +static void rt_load_method_maybe(mp_obj_t base, qstr attr, mp_obj_t *dest) { // clear output to indicate no attribute/method found yet dest[0] = MP_OBJ_NULL; dest[1] = MP_OBJ_NULL; @@ -830,7 +832,7 @@ void rt_load_method(mp_obj_t base, qstr attr, mp_obj_t *dest) { } // if nothing found yet, look for built-in and generic names - if (dest[0] == NULL) { + if (dest[0] == MP_OBJ_NULL) { if (attr == MP_QSTR___next__ && type->iternext != NULL) { dest[0] = (mp_obj_t)&mp_builtin_next_obj; dest[1] = base; @@ -861,8 +863,14 @@ void rt_load_method(mp_obj_t base, qstr attr, mp_obj_t *dest) { } } } +} - if (dest[0] == NULL) { +void rt_load_method(mp_obj_t base, qstr attr, mp_obj_t *dest) { + DEBUG_OP_printf("load method %p.%s\n", base, qstr_str(attr)); + + rt_load_method_maybe(base, attr, dest); + + if (dest[0] == MP_OBJ_NULL) { // no attribute/method called attr // following CPython, we give a more detailed error message for type objects if (MP_OBJ_IS_TYPE(base, &mp_const_type)) { @@ -910,7 +918,16 @@ mp_obj_t rt_getiter(mp_obj_t o_in) { if (type->getiter != NULL) { return type->getiter(o_in); } else { - nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_TypeError, "'%s' object is not iterable", type->name)); + // check for __getitem__ method + mp_obj_t dest[2]; + rt_load_method_maybe(o_in, qstr_from_str("__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_QSTR_TypeError, "'%s' object is not iterable", type->name)); + } } } diff --git a/tests/basics/getitem.py b/tests/basics/getitem.py new file mode 100644 index 0000000000..f39296aca3 --- /dev/null +++ b/tests/basics/getitem.py @@ -0,0 +1,22 @@ +# create a class that has a __getitem__ method +class A: + def __getitem__(self, index): + print('getitem', index) + if index > 10: + raise StopIteration + +# test __getitem__ +A()[0] +A()[1] + +# iterate using a for loop +for i in A(): + pass + +# iterate manually +it = iter(A()) +try: + while True: + next(it) +except StopIteration: + pass