unix/modos: Implement ilistdir().

pfalcon · pfalcon · commit d874702fe1e5 · 2015-12-14T23:48:12.000+02:00
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().
diff --git a/unix/modos.c b/unix/modos.c
@@ -31,6 +31,7 @@
 #include <errno.h>
 #include <stdlib.h>
 #include <string.h>
+#include <dirent.h>
 #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);
diff --git a/unix/mpconfigport.h b/unix/mpconfigport.h
@@ -250,3 +250,12 @@ extern const struct _mp_obj_fun_builtin_t mp_builtin_open_obj;
 #include <alloca.h>
 #endif
 #endif
+
+// From "man readdir": "Under glibc, programs can check for the availability
+// of the fields [in struct dirent] not defined in POSIX.1 by testing whether
+// the macros [...], _DIRENT_HAVE_D_TYPE are defined."
+// Other libc's don't define it, but proactively assume that dirent->d_type
+// is available on a modern *nix system.
+#ifndef _DIRENT_HAVE_D_TYPE
+#define _DIRENT_HAVE_D_TYPE (1)
+#endif
diff --git a/unix/qstrdefsport.h b/unix/qstrdefsport.h
@@ -43,6 +43,7 @@ Q(system)
 Q(unlink)
 Q(getenv)
 Q(mkdir)
+Q(ilistdir)
 
 Q(uselect)
 Q(poll)
