py/ringbuf: Add micropython.ringbuffer() interface for general use.

pi-anl · pi-anl · commit e81a9ba2aac0 · 2024-08-07T15:56:53.000+10:00
Signed-off-by: Andrew Leech &lt;andrew.leech@planetinnovation.com.au&gt;
diff --git a/docs/library/micropython.rst b/docs/library/micropython.rst
@@ -147,3 +147,15 @@ Functions
 
    There is a finite queue to hold the scheduled functions and `schedule()`
    will raise a `RuntimeError` if the queue is full.
+
+.. function:: ringbuffer(size)
+.. function:: ringbuffer(buffer)
+   :noindex:
+
+   Provides a fixed-size ringbuffer with stream interface. Can be used similar to
+   `io.BytesIO` however has separate read/write positions.
+
+   Can be created with integer size provided and a suitable buffer will be created,
+   after which no further allocations will take place during use.
+   Alternatively a `bytearray`, `memoryview` or similar object can be provided at
+   init for in-place use.
diff --git a/py/modmicropython.c b/py/modmicropython.c
@@ -31,6 +31,7 @@
 #include "py/runtime.h"
 #include "py/gc.h"
 #include "py/mphal.h"
+#include "py/ringbuf.h"
 
 #if MICROPY_PY_MICROPYTHON
 
@@ -203,6 +204,9 @@ static const mp_rom_map_elem_t mp_module_micropython_globals_table[] = {
     #if MICROPY_ENABLE_SCHEDULER
     { MP_ROM_QSTR(MP_QSTR_schedule), MP_ROM_PTR(&mp_micropython_schedule_obj) },
     #endif
+    #if MICROPY_PY_MICROPYTHON_RINGBUFFER
+    { MP_ROM_QSTR(MP_QSTR_ringbuffer), MP_ROM_PTR(&mp_type_micropython_ringbuffer) },
+    #endif
 };
 
 static MP_DEFINE_CONST_DICT(mp_module_micropython_globals, mp_module_micropython_globals_table);
diff --git a/py/mpconfig.h b/py/mpconfig.h
@@ -981,6 +981,11 @@ typedef double mp_float_t;
 #define MICROPY_ENABLE_SCHEDULER (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)
 #endif
 
+// Support for micropython.ringbuffer()
+#ifndef MICROPY_PY_MICROPYTHON_RINGBUFFER
+#define MICROPY_PY_MICROPYTHON_RINGBUFFER (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)
+#endif
+
 // Whether the scheduler supports scheduling static nodes with C callbacks
 #ifndef MICROPY_SCHEDULER_STATIC_NODES
 #define MICROPY_SCHEDULER_STATIC_NODES (0)
diff --git a/py/ringbuf.c b/py/ringbuf.c
@@ -118,3 +118,175 @@ int ringbuf_put_bytes(ringbuf_t *r, const uint8_t *data, size_t data_len) {
     r->iput = iput_a;
     return 0;
 }
+
+#if MICROPY_PY_MICROPYTHON_RINGBUFFER
+
+#include "py/runtime.h"
+#include "py/stream.h"
+#include "py/mphal.h"
+
+typedef struct _micropython_ringbuffer_obj_t {
+    mp_obj_base_t base;
+    ringbuf_t ringbuffer;
+    mp_int_t timeout;       // timeout waiting for first char (in ms)
+} micropython_ringbuffer_obj_t;
+
+static mp_obj_t micropython_ringbuffer_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
+    mp_arg_check_num(n_args, n_kw, 1, 2, false);
+    mp_int_t buff_size = -1;
+    mp_buffer_info_t bufinfo = {NULL, 0, 0};
+
+    if (!mp_get_buffer(args[0], &bufinfo, MP_BUFFER_RW)) {
+        buff_size = mp_obj_get_int(args[0]);
+    }
+    micropython_ringbuffer_obj_t *self = mp_obj_malloc(micropython_ringbuffer_obj_t, type);
+    if (bufinfo.buf != NULL) {
+        // buffer passed in, use it directly for ringbuffer
+        self->ringbuffer.buf = bufinfo.buf;
+        self->ringbuffer.size = bufinfo.len;
+        self->ringbuffer.iget = self->ringbuffer.iput = 0;
+    } else {
+        // Allocation buffer, add one extra to buff_size as ringbuf consumes one byte for tracking.
+        ringbuf_alloc(&(self->ringbuffer), buff_size + 1);
+    }
+
+    if (n_args > 1) {
+        self->timeout = mp_obj_get_int(args[1]);
+    } else {
+        self->timeout = -1;
+    }
+    return MP_OBJ_FROM_PTR(self);
+}
+
+static mp_obj_t micropython_ringbuffer_settimeout(mp_obj_t self_in, mp_obj_t timeout_in) {
+    micropython_ringbuffer_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    self->timeout = mp_obj_get_int(timeout_in);
+    return mp_const_none;
+}
+static MP_DEFINE_CONST_FUN_OBJ_2(micropython_ringbuffer_settimeout_obj, micropython_ringbuffer_settimeout);
+
+
+static mp_uint_t micropython_ringbuffer_read(mp_obj_t self_in, void *buf_in, mp_uint_t size, int *errcode) {
+    micropython_ringbuffer_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    mp_uint_t start = mp_hal_ticks_ms();
+
+    size = MIN(size, ((mp_uint_t)self->ringbuffer.size - 1)); // limit size to ringbuffer length
+
+    while (ringbuf_get_bytes(&self->ringbuffer, buf_in, size) == -1) {
+        if (self->timeout > -1 && (mp_hal_ticks_ms() - start > (mp_uint_t)self->timeout)) {
+            // timed out
+            if ((size = MIN(size, ringbuf_avail(&self->ringbuffer))) > 0) {
+                // return available data
+                ringbuf_get_bytes(&self->ringbuffer, buf_in, size);
+                *errcode = 0;
+                return size;
+            }
+            // no data available
+            *errcode = MP_EAGAIN;
+            return MP_STREAM_ERROR;
+        }
+        mp_event_handle_nowait();
+    }
+    *errcode = 0;
+    return size;
+}
+
+static mp_uint_t micropython_ringbuffer_write(mp_obj_t self_in, const void *buf_in, mp_uint_t size, int *errcode) {
+    micropython_ringbuffer_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    mp_uint_t start = mp_hal_ticks_ms();
+
+    size = MIN(size, ((mp_uint_t)self->ringbuffer.size - 1)); // limit size to ringbuffer length
+
+    while (ringbuf_put_bytes(&self->ringbuffer, buf_in, size) == -1) {
+        if (self->timeout > -1 && (mp_hal_ticks_ms() - start > (mp_uint_t)self->timeout)) {
+            // timed out
+            if ((size = MIN(size, ringbuf_free(&self->ringbuffer))) > 0) {
+                // write whatever can fit
+                ringbuf_put_bytes(&self->ringbuffer, buf_in, size);
+                *errcode = 0;
+                return size;
+            }
+            // no space available
+            *errcode = MP_EAGAIN;
+            return MP_STREAM_ERROR;
+        }
+        mp_event_handle_nowait();
+    }
+    *errcode = 0;
+    return size;
+
+}
+
+static mp_uint_t micropython_ringbuffer_ioctl(mp_obj_t self_in, mp_uint_t request, uintptr_t arg, int *errcode) {
+    micropython_ringbuffer_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    mp_uint_t ret;
+    if (request == MP_STREAM_POLL) {
+        ret = 0;
+        if ((arg & MP_STREAM_POLL_RD) && ringbuf_avail(&self->ringbuffer) > 0) {
+            ret |= MP_STREAM_POLL_RD;
+        }
+        if ((arg & MP_STREAM_POLL_WR) && ringbuf_free(&self->ringbuffer) > 0) {
+            ret |= MP_STREAM_POLL_WR;
+        }
+    } else if (request == MP_STREAM_FLUSH) {
+        // Should we wait here until empty / timeout?
+        ret = 0;
+    } else if (request == MP_STREAM_CLOSE) {
+        // We don't want to reset head/tail pointers as there might
+        // still be someone using it, eg. if ringbuffer is used instead of
+        // a socket, a "writer" might call close before the "reader" is
+        // finished.
+        // Should we flush here though?
+        ret = 0;
+    } else {
+        *errcode = MP_EINVAL;
+        ret = MP_STREAM_ERROR;
+    }
+    return ret;
+}
+
+static mp_obj_t micropython_ringbuffer_any(mp_obj_t self_in) {
+    micropython_ringbuffer_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    return MP_OBJ_NEW_SMALL_INT(ringbuf_avail(&self->ringbuffer));
+}
+static MP_DEFINE_CONST_FUN_OBJ_1(micropython_ringbuffer_any_obj, micropython_ringbuffer_any);
+
+static mp_obj_t micropython_ringbuffer_reset(mp_obj_t self_in) {
+    micropython_ringbuffer_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    self->ringbuffer.iget = self->ringbuffer.iput = 0;
+    return mp_const_none;
+}
+static MP_DEFINE_CONST_FUN_OBJ_1(micropython_ringbuffer_reset_obj, micropython_ringbuffer_reset);
+
+
+static const mp_rom_map_elem_t micropython_ringbuffer_locals_dict_table[] = {
+    { MP_ROM_QSTR(MP_QSTR_any), MP_ROM_PTR(&micropython_ringbuffer_any_obj) },
+    { MP_ROM_QSTR(MP_QSTR_settimeout), MP_ROM_PTR(&micropython_ringbuffer_settimeout_obj) },
+    { MP_ROM_QSTR(MP_QSTR_reset), MP_ROM_PTR(&micropython_ringbuffer_reset_obj) },
+    { MP_ROM_QSTR(MP_QSTR_flush), MP_ROM_PTR(&mp_stream_flush_obj) },
+    { MP_ROM_QSTR(MP_QSTR_read), MP_ROM_PTR(&mp_stream_read_obj) },
+    { MP_ROM_QSTR(MP_QSTR_readline), MP_ROM_PTR(&mp_stream_unbuffered_readline_obj) },
+    { MP_ROM_QSTR(MP_QSTR_readinto), MP_ROM_PTR(&mp_stream_readinto_obj) },
+    { MP_ROM_QSTR(MP_QSTR_write), MP_ROM_PTR(&mp_stream_write_obj) },
+    { MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&mp_stream_close_obj) },
+
+};
+static MP_DEFINE_CONST_DICT(micropython_ringbuffer_locals_dict, micropython_ringbuffer_locals_dict_table);
+
+static const mp_stream_p_t ringbuffer_stream_p = {
+    .read = micropython_ringbuffer_read,
+    .write = micropython_ringbuffer_write,
+    .ioctl = micropython_ringbuffer_ioctl,
+    .is_text = false,
+};
+
+MP_DEFINE_CONST_OBJ_TYPE(
+    mp_type_micropython_ringbuffer,
+    MP_QSTR_ringbuffer,
+    MP_TYPE_FLAG_NONE,
+    make_new, micropython_ringbuffer_make_new,
+    protocol, &ringbuffer_stream_p,
+    locals_dict, &micropython_ringbuffer_locals_dict
+    );
+
+#endif
diff --git a/py/ringbuf.h b/py/ringbuf.h
@@ -28,9 +28,10 @@
 
 #include <stddef.h>
 #include <stdint.h>
+#include "py/mpconfig.h"
 
-#ifdef _MSC_VER
-#include "py/mpconfig.h" // For inline.
+#if MICROPY_PY_MICROPYTHON_RINGBUFFER
+#include "py/obj.h"
 #endif
 
 typedef struct _ringbuf_t {
@@ -99,4 +100,8 @@ int ringbuf_put16(ringbuf_t *r, uint16_t v);
 int ringbuf_get_bytes(ringbuf_t *r, uint8_t *data, size_t data_len);
 int ringbuf_put_bytes(ringbuf_t *r, const uint8_t *data, size_t data_len);
 
+#if MICROPY_PY_MICROPYTHON_RINGBUFFER
+extern const mp_obj_type_t mp_type_micropython_ringbuffer;
+#endif
+
 #endif // MICROPY_INCLUDED_PY_RINGBUF_H
diff --git a/tests/micropython/ringbuffer.py b/tests/micropython/ringbuffer.py
@@ -0,0 +1,39 @@
+# check that micropython.ringbuffer works correctly.
+
+import micropython
+
+try:
+    micropython.ringbuffer
+except AttributeError:
+    print("SKIP")
+    raise SystemExit
+
+rb = micropython.ringbuffer(16, 1)
+print(rb)
+
+print(rb.any())
+
+rb.write(b"\x00")
+print(rb.any())
+
+rb.write(b"\x00")
+print(rb.any())
+
+print(rb.read(2))
+print(rb.any())
+
+
+rb.write(b"\x00\x01")
+print(rb.read())
+
+print(rb.read(1))
+
+print(rb.write(b"\x00\x01" * 10))
+print(rb.read())
+print(rb.read())
+
+try:
+    # size must be int.
+    micropython.ringbuffer(None)
+except TypeError as ex:
+    print(ex)
diff --git a/tests/micropython/ringbuffer.py.exp b/tests/micropython/ringbuffer.py.exp
@@ -0,0 +1,11 @@
+<ringbuffer>
+0
+1
+2
+b'\x00\x00'
+0
+b'\x00\x01'
+None
+16
+b'\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01'
+can't convert NoneType to int
diff --git a/tests/micropython/ringbuffer_async.py b/tests/micropython/ringbuffer_async.py
@@ -0,0 +1,36 @@
+# check that micropython.ringbuffer works correctly with asyncio.Stream.
+
+import micropython
+
+try:
+    import asyncio
+
+    micropython.ringbuffer
+except AttributeError:
+    print("SKIP")
+    raise SystemExit
+
+rb = micropython.ringbuffer(16, 0)
+rba = asyncio.StreamWriter(rb)
+
+data = b"ABC123" * 20
+print("w", len(data))
+
+
+async def data_writer():
+    global data
+    rba.write(data)
+    await rba.drain()
+
+
+async def main():
+    task = asyncio.create_task(data_writer())
+    await asyncio.sleep_ms(10)
+    # buff = bytearray(len(data))
+    read = await rba.readexactly(len(data))
+    print(read)
+    print("r", len(read))
+    print(read == data)
+
+
+asyncio.run(main())
diff --git a/tests/micropython/ringbuffer_async.py.exp b/tests/micropython/ringbuffer_async.py.exp
@@ -0,0 +1,4 @@
+w 120
+b'ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123'
+r 120
+True

-Original file line number
+Diff line change
@@ @@ -0,0 +1,4 @@ @@
 +w 120
 +b'ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123'
 +r 120
 +True