py/modio: Implement BufferedReader

klondi · klondi · commit 3ce4d08ce80c · 2024-01-09T09:01:40.000+01:00
Some times there is a need for a BufferedReader in a way similar to how
BufferedWritter works. A clear example is when using an underlying device
requiring aligned reads, but a less clear example is when using
deflate.DeflateIO which will do only 1-byte reads and can become
crippling quickly when the underlying object is a python implemented
stream instead of a native one.

The BufferedReader will only attempt to do full-buffer reads and
ensures word-alignment in a way similar to how the writer does.
Similarly, it will also hide any errors when partial reads happen to
ensure that any data copied so far can be returned first.

Signed-off-by: Francisco Blas (klondike) Izquierdo Riera &lt;klondike@klondike.es&gt;
diff --git a/ports/esp32/mpconfigport.h b/ports/esp32/mpconfigport.h
@@ -77,6 +77,7 @@
 #define MICROPY_PY_ALL_INPLACE_SPECIAL_METHODS (1)
 #define MICROPY_PY_BUILTINS_HELP_TEXT       esp32_help_text
 #define MICROPY_PY_IO_BUFFEREDWRITER        (1)
+#define MICROPY_PY_IO_BUFFEREDREADER        (1)
 #define MICROPY_PY_TIME_GMTIME_LOCALTIME_MKTIME (1)
 #define MICROPY_PY_TIME_TIME_TIME_NS        (1)
 #define MICROPY_PY_TIME_INCLUDEFILE         "ports/esp32/modtime.c"
diff --git a/py/modio.c b/py/modio.c
@@ -108,7 +108,8 @@ STATIC MP_DEFINE_CONST_OBJ_TYPE(
 
 #endif // MICROPY_PY_IO_IOBASE
 
-#if MICROPY_PY_IO_BUFFEREDWRITER
+#if MICROPY_PY_IO_BUFFEREDWRITER || MICROPY_PY_IO_BUFFEREDREADER
+// The structure and the new method are shared by reader and writer
 typedef struct _mp_obj_bufwriter_t {
     mp_obj_base_t base;
     mp_obj_t stream;
@@ -129,6 +130,9 @@ STATIC mp_obj_t bufwriter_make_new(const mp_obj_type_t *type, size_t n_args, siz
     return o;
 }
 
+#endif // MICROPY_PY_IO_BUFFEREDWRITER || MICROPY_PY_IO_BUFFEREDREADER
+
+#if MICROPY_PY_IO_BUFFEREDWRITER
 // Writes out the data stored in the buffer so far
 STATIC int bufwriter_do_write(mp_obj_bufwriter_t *self) {
     int rv = 0;
@@ -240,6 +244,84 @@ STATIC MP_DEFINE_CONST_OBJ_TYPE(
     );
 #endif // MICROPY_PY_IO_BUFFEREDWRITER
 
+#if MICROPY_PY_IO_BUFFEREDREADER
+STATIC mp_uint_t bufreader_read(mp_obj_t self_in, void *buf, mp_uint_t size, int *errcode) {
+    mp_obj_bufwriter_t *self = MP_OBJ_TO_PTR(self_in);
+
+    mp_uint_t org_size = size;
+    // Cache this value since it should not change thus avoiding a dereference
+    size_t alloc = self->alloc;
+
+    *errcode = self->error;
+    self->error = 0;
+
+    while (*errcode == 0 && size > 0) {
+        // Buffer filling policy here is to fill the entire buffer all the time.
+        // This allows e.g. to have a block device as backing storage and read
+        // entire blocks from it. memcpy below is not ideal and could be
+        // optimized in some cases. But the way it is now it at least ensures
+        // that buffer is word-aligned, to guard against obscure cases when it
+        // matters, e.g.
+        // https://github.com/micropython/micropython/issues/1863
+
+        // Buffer needs to have at least one byte
+        if (self->len == 0) {
+            self->len = mp_stream_read_exactly(self->stream, self->buf, alloc, errcode);
+            if (self->len < alloc) {
+                // If no data was read stop the loop as it is either an error or an EOF
+                // This check is moved here because it speeds up the more common case
+                // where the buffer was completely filled.
+                if (self->len == 0) {
+                    break;
+                }
+                // Buffers may overlap, therefore we use memmove instead of memcpy
+                memmove(self->buf + (alloc - self->len), self->buf, self->len);
+            }
+        }
+        mp_uint_t rem = MIN(self->len, size);
+        memcpy(buf, self->buf + (alloc - self->len), rem);
+        self->len -= rem;
+        buf = (byte *)buf + rem;
+        size -= rem;
+    }
+
+    // If there is a prior error
+    if (*errcode != 0) {
+        // If no data was read return it
+        if (size == org_size) {
+            return MP_STREAM_ERROR;
+        }
+        // Otherwise save it so we can fail on the next call
+        self->error = *errcode;
+        *errcode = 0;
+    }
+
+    // Return the data that has been read so far
+    return org_size - size;
+}
+
+STATIC const mp_rom_map_elem_t bufreader_locals_dict_table[] = {
+    { MP_ROM_QSTR(MP_QSTR_read), MP_ROM_PTR(&mp_stream_read_obj) },
+    { MP_ROM_QSTR(MP_QSTR_readinto), MP_ROM_PTR(&mp_stream_readinto_obj) },
+    { MP_ROM_QSTR(MP_QSTR_readline), MP_ROM_PTR(&mp_stream_unbuffered_readline_obj) },
+};
+STATIC MP_DEFINE_CONST_DICT(bufreader_locals_dict, bufreader_locals_dict_table);
+
+STATIC const mp_stream_p_t bufreader_stream_p = {
+    .read = bufreader_read,
+};
+
+STATIC MP_DEFINE_CONST_OBJ_TYPE(
+    mp_type_bufreader,
+    MP_QSTR_BufferedReader,
+    MP_TYPE_FLAG_NONE,
+    make_new, bufwriter_make_new,
+    protocol, &bufreader_stream_p,
+    locals_dict, &bufreader_locals_dict
+    );
+#endif // MICROPY_PY_IO_BUFFEREDREADER
+
+
 STATIC const mp_rom_map_elem_t mp_module_io_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_io) },
     // Note: mp_builtin_open_obj should be defined by port, it's not
@@ -255,6 +337,9 @@ STATIC const mp_rom_map_elem_t mp_module_io_globals_table[] = {
     #if MICROPY_PY_IO_BUFFEREDWRITER
     { MP_ROM_QSTR(MP_QSTR_BufferedWriter), MP_ROM_PTR(&mp_type_bufwriter) },
     #endif
+    #if MICROPY_PY_IO_BUFFEREDREADER
+    { MP_ROM_QSTR(MP_QSTR_BufferedReader), MP_ROM_PTR(&mp_type_bufreader) },
+    #endif
 };
 
 STATIC MP_DEFINE_CONST_DICT(mp_module_io_globals, mp_module_io_globals_table);
diff --git a/py/mpconfig.h b/py/mpconfig.h
@@ -1396,6 +1396,11 @@ typedef double mp_float_t;
 #define MICROPY_PY_IO_BUFFEREDWRITER (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EVERYTHING)
 #endif
 
+// Whether to provide "io.BufferedReader" class
+#ifndef MICROPY_PY_IO_BUFFEREDREADER
+#define MICROPY_PY_IO_BUFFEREDREADER (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EVERYTHING)
+#endif
+
 // Whether to provide "struct" module
 #ifndef MICROPY_PY_STRUCT
 #define MICROPY_PY_STRUCT (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_CORE_FEATURES)
