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

6 files changed, 58 insertions, 29 deletions

diff --git a/py/obj.h b/py/obj.h
index b52b5a0c7e..dd728c9c2f 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -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);
diff --git a/py/objexcept.c b/py/objexcept.c
index 0650920e7d..79b85b28b7 100644
--- a/py/objexcept.c
+++ b/py/objexcept.c
@@ -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) {
diff --git a/py/objtype.c b/py/objtype.c
index 4898c5ccb4..9cb29744c7 100644
--- a/py/objtype.c
+++ b/py/objtype.c
@@ -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) {
diff --git a/py/qstrdefs.h b/py/qstrdefs.h
index 44c61002a6..1c8afe797d 100644
--- a/py/qstrdefs.h
+++ b/py/qstrdefs.h
@@ -17,6 +17,7 @@ Q(__repl_print__)
 
 Q(__bool__)
 Q(__len__)
+Q(__iter__)
 Q(__getitem__)
 Q(__setitem__)
 Q(__add__)
diff --git a/py/runtime.c b/py/runtime.c
index 0f5d1229f2..3e0ef98d8b 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -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)));
+        }
     }
 }
 
diff --git a/py/vm.c b/py/vm.c
index 22243c403e..8ae619f9be 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -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.