From d874702fe1e5822f8a3e1662657ac27d0b15bb68 Mon Sep 17 00:00:00 2001 From: Paul Sokolovsky Date: Mon, 14 Dec 2015 23:48:12 +0200 Subject: unix/modos: Implement ilistdir(). ilistdir() returns iterator which yields triples of (name, type, ino) where ino is inode number for entry's data, type of entry (file/dir/etc.), and name of file/dir. listdir() can be easily implemented in terms of this iterator (which is otherwise more efficient in terms of memory use and may save expensive call to stat() for each returned entry). CPython has os.scandir() which also returns an iterator, but it yields more complex objects of DirEntry type. scandir() can also be easily implemented in terms of ilistdir(). --- unix/modos.c | 47 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) (limited to 'unix/modos.c') diff --git a/unix/modos.c b/unix/modos.c index f681acaff2..abc96e638d 100644 --- a/unix/modos.c +++ b/unix/modos.c @@ -31,6 +31,7 @@ #include #include #include +#include #include "py/mpconfig.h" #include "py/nlr.h" @@ -155,6 +156,51 @@ STATIC mp_obj_t mod_os_mkdir(mp_obj_t path_in) { } STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_os_mkdir_obj, mod_os_mkdir); +typedef struct _mp_obj_listdir_t { + mp_obj_base_t base; + mp_fun_1_t iternext; + DIR *dir; +} mp_obj_listdir_t; + +STATIC mp_obj_t listdir_next(mp_obj_t self_in) { + mp_obj_listdir_t *self = MP_OBJ_TO_PTR(self_in); + + if (self->dir == NULL) { + goto done; + } + struct dirent *dirent = readdir(self->dir); + if (dirent == NULL) { + closedir(self->dir); + self->dir = NULL; + done: + return MP_OBJ_STOP_ITERATION; + } + + mp_obj_tuple_t *t = MP_OBJ_TO_PTR(mp_obj_new_tuple(3, NULL)); + t->items[0] = mp_obj_new_str(dirent->d_name, strlen(dirent->d_name), false); + #ifdef _DIRENT_HAVE_D_TYPE + t->items[1] = MP_OBJ_NEW_SMALL_INT(dirent->d_type); + #else + // DT_UNKNOWN should have 0 value on any reasonable system + t->items[1] = 0; + #endif + t->items[2] = MP_OBJ_NEW_SMALL_INT(dirent->d_ino); + return MP_OBJ_FROM_PTR(t); +} + +STATIC mp_obj_t mod_os_ilistdir(mp_uint_t n_args, const mp_obj_t *args) { + const char *path = "."; + if (n_args > 0) { + path = mp_obj_str_get_str(args[0]); + } + mp_obj_listdir_t *o = m_new_obj(mp_obj_listdir_t); + o->base.type = &mp_type_polymorph_iter; + o->dir = opendir(path); + o->iternext = listdir_next; + return MP_OBJ_FROM_PTR(o); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_os_ilistdir_obj, 0, 1, mod_os_ilistdir); + STATIC const mp_rom_map_elem_t mp_module_os_globals_table[] = { { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_uos) }, { MP_ROM_QSTR(MP_QSTR_stat), MP_ROM_PTR(&mod_os_stat_obj) }, @@ -165,6 +211,7 @@ STATIC const mp_rom_map_elem_t mp_module_os_globals_table[] = { { MP_ROM_QSTR(MP_QSTR_unlink), MP_ROM_PTR(&mod_os_unlink_obj) }, { MP_ROM_QSTR(MP_QSTR_getenv), MP_ROM_PTR(&mod_os_getenv_obj) }, { MP_ROM_QSTR(MP_QSTR_mkdir), MP_ROM_PTR(&mod_os_mkdir_obj) }, + { MP_ROM_QSTR(MP_QSTR_ilistdir), MP_ROM_PTR(&mod_os_ilistdir_obj) }, }; STATIC MP_DEFINE_CONST_DICT(mp_module_os_globals, mp_module_os_globals_table); -- cgit v1.2.3