From d874702fe1e5822f8a3e1662657ac27d0b15bb68 Mon Sep 17 00:00:00 2001
From: Paul Sokolovsky <pfalcon@users.sourceforge.net>
Date: Mon, 14 Dec 2015 23:48:12 +0200
Subject: [PATCH] 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 +++++++++++++++++++++++++++++++++++++++++++++
 unix/mpconfigport.h |  9 +++++++++
 unix/qstrdefsport.h |  1 +
 3 files changed, 57 insertions(+)

diff --git a/unix/modos.c b/unix/modos.c
index f681acaff..abc96e638 100644
--- 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
index ad75f5287..5fd122e96 100644
--- 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
index 608234e23..9965c030c 100644
--- 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)
-- 
GitLab