shared-module: Fix os.listdir to correctly return list rather than

tannewt · tannewt · commit c797f3cb240e · 2017-06-28T12:51:32.000-07:00
iterator.
diff --git a/shared-module/os/__init__.c b/shared-module/os/__init__.c
@@ -44,7 +44,7 @@ STATIC mp_vfs_mount_t *lookup_path(const char* path, mp_obj_t *path_out) {
     const char *p_out;
     mp_vfs_mount_t *vfs = mp_vfs_lookup_path(path, &p_out);
     if (vfs != MP_VFS_NONE && vfs != MP_VFS_ROOT) {
-        *path_out = mp_obj_new_str_of_type(mp_obj_get_type(path),
+        *path_out = mp_obj_new_str_of_type(&mp_type_str,
             (const byte*)p_out, strlen(p_out));
     }
     return vfs;
@@ -95,18 +95,28 @@ mp_obj_t common_hal_os_listdir(const char* path) {
     mp_obj_t path_out;
     mp_vfs_mount_t *vfs = lookup_path(path, &path_out);
 
+    mp_vfs_ilistdir_it_t iter;
+    mp_obj_t iter_obj = MP_OBJ_FROM_PTR(&iter);
+
     if (vfs == MP_VFS_ROOT) {
         // list the root directory
-        mp_vfs_ilistdir_it_t *iter = m_new_obj(mp_vfs_ilistdir_it_t);
-        iter->base.type = &mp_type_polymorph_iter;
-        iter->iternext = mp_vfs_ilistdir_it_iternext;
-        iter->cur.vfs = MP_STATE_VM(vfs_mount_table);
-        iter->is_str = mp_obj_get_type(path) == &mp_type_str;
-        iter->is_iter = false;
-        return MP_OBJ_FROM_PTR(iter);
+        iter.base.type = &mp_type_polymorph_iter;
+        iter.iternext = mp_vfs_ilistdir_it_iternext;
+        iter.cur.vfs = MP_STATE_VM(vfs_mount_table);
+        iter.is_str = mp_obj_get_type(path) == &mp_type_str;
+        iter.is_iter = false;
+    } else {
+        iter_obj = mp_vfs_proxy_call(vfs, MP_QSTR_ilistdir, 1, &path_out);
     }
 
-    return mp_vfs_proxy_call(vfs, MP_QSTR_ilistdir, 1, &path_out);
+    mp_obj_t dir_list = mp_obj_new_list(0, NULL);
+    mp_obj_t next;
+    while ((next = mp_iternext(iter_obj)) != MP_OBJ_STOP_ITERATION) {
+        mp_obj_t *items;
+        mp_obj_get_array_fixed_n(next, 3, &items);
+        mp_obj_list_append(dir_list, items[0]);
+    }
+    return dir_list;
 }
 
 void common_hal_os_mkdir(const char* path) {
