1 files changed, 65 insertions, 0 deletions

diff --git a/py/objcode.c b/py/objcode.c
index 9b98a69679..52df84d012 100644
--- a/py/objcode.c
+++ b/py/objcode.c
@@ -107,6 +107,67 @@ static mp_obj_t raw_code_lnotab(const mp_raw_code_t *rc) {
     return o;
 }
 
+static mp_obj_t code_colines_iter(mp_obj_t);
+static mp_obj_t code_colines_next(mp_obj_t);
+typedef struct _mp_obj_colines_iter_t {
+    mp_obj_base_t base;
+    mp_fun_1_t iternext;
+    const mp_raw_code_t *rc;
+    mp_uint_t bc;
+    mp_uint_t source_line;
+    const byte *ci;
+} mp_obj_colines_iter_t;
+
+static mp_obj_t code_colines_iter(mp_obj_t self_in) {
+    mp_obj_code_t *self = MP_OBJ_TO_PTR(self_in);
+    mp_obj_colines_iter_t *iter = mp_obj_malloc(mp_obj_colines_iter_t, &mp_type_polymorph_iter);
+    iter->iternext = code_colines_next;
+    iter->rc = self->rc;
+    iter->bc = 0;
+    iter->source_line = 1;
+    iter->ci = self->rc->prelude.line_info;
+    return MP_OBJ_FROM_PTR(iter);
+}
+static MP_DEFINE_CONST_FUN_OBJ_1(code_colines_obj, code_colines_iter);
+
+static mp_obj_t code_colines_next(mp_obj_t iter_in) {
+    mp_obj_colines_iter_t *iter = MP_OBJ_TO_PTR(iter_in);
+    const byte *ci_end = iter->rc->prelude.line_info_top;
+
+    mp_uint_t start = iter->bc;
+    mp_uint_t line_no = iter->source_line;
+    bool another = true;
+
+    while (another && iter->ci < ci_end) {
+        another = false;
+        mp_code_lineinfo_t decoded = mp_bytecode_decode_lineinfo(&iter->ci);
+        iter->bc += decoded.bc_increment;
+        iter->source_line += decoded.line_increment;
+
+        if (decoded.bc_increment == 0) {
+            line_no = iter->source_line;
+            another = true;
+        } else if (decoded.line_increment == 0) {
+            another = true;
+        }
+    }
+
+    if (another) {
+        mp_uint_t prelude_size = (iter->rc->prelude.opcodes - (const byte *)iter->rc->fun_data);
+        mp_uint_t bc_end = iter->rc->fun_data_len - prelude_size;
+        if (iter->bc >= bc_end) {
+            return MP_OBJ_STOP_ITERATION;
+        } else {
+            iter->bc = bc_end;
+        }
+    }
+
+    mp_uint_t end = iter->bc;
+    mp_obj_t next[3] = {MP_OBJ_NEW_SMALL_INT(start), MP_OBJ_NEW_SMALL_INT(end), MP_OBJ_NEW_SMALL_INT(line_no)};
+
+    return mp_obj_new_tuple(MP_ARRAY_SIZE(next), next);
+}
+
 static void code_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
     if (dest[0] != MP_OBJ_NULL) {
         // not load attribute
@@ -143,6 +204,10 @@ static void code_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
             }
             dest[0] = o->lnotab;
             break;
+        case MP_QSTR_co_lines:
+            dest[0] = MP_OBJ_FROM_PTR(&code_colines_obj);
+            dest[1] = self_in;
+            break;
     }
 }