From 957f20a31bd957267b4d2e8bd1a7cebca882c0c5 Mon Sep 17 00:00:00 2001 From: Andrew Leech Date: Mon, 26 Sep 2022 11:02:31 +1000 Subject: [PATCH] py/ringbuf: Add micropython.ringbuffer() interface for general use. Signed-off-by: Andrew Leech --- docs/library/micropython.rst | 12 +++ py/modmicropython.c | 4 + py/mpconfig.h | 5 + py/ringbuf.c | 125 ++++++++++++++++++++++ py/ringbuf.h | 9 +- tests/micropython/ringbuffer.py | 38 +++++++ tests/micropython/ringbuffer.py.exp | 11 ++ tests/micropython/ringbuffer_async.py | 36 +++++++ tests/micropython/ringbuffer_async.py.exp | 4 + 9 files changed, 242 insertions(+), 2 deletions(-) create mode 100644 tests/micropython/ringbuffer.py create mode 100644 tests/micropython/ringbuffer.py.exp create mode 100644 tests/micropython/ringbuffer_async.py create mode 100644 tests/micropython/ringbuffer_async.py.exp diff --git a/docs/library/micropython.rst b/docs/library/micropython.rst index b17dfa9a75a48..00bc4bdf3e6c1 100644 --- 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 index af6ad01795f67..354147fd17804 100644 --- 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 index 98893ceb6d97b..8c84201a748d2 100644 --- 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 index 10dca62081e8c..9cecdb4a1f5e7 100644 --- a/py/ringbuf.c +++ b/py/ringbuf.c @@ -118,3 +118,128 @@ 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; +} 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, 1, 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. + // This can be user to no-copy stream an existing data buffer (except final byte). + self->ringbuffer.buf = bufinfo.buf; + self->ringbuffer.size = bufinfo.len; + self->ringbuffer.iput = bufinfo.len - 1; + self->ringbuffer.iget = 0; + } else { + // Allocate new buffer, add one extra to buff_size as ringbuf consumes one byte for tracking. + ringbuf_alloc(&(self->ringbuffer), buff_size + 1); + } + return MP_OBJ_FROM_PTR(self); +} + +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); + size = MIN(size, ringbuf_avail(&self->ringbuffer)); // limit size to available data + + if (size == 0 || ringbuf_get_bytes(&self->ringbuffer, buf_in, size) == -1) { + // no data available + *errcode = MP_EAGAIN; + return MP_STREAM_ERROR; + } + *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); + size = MIN(size, ringbuf_free(&self->ringbuffer)); // limit size to available space + + if (size == 0 || ringbuf_put_bytes(&self->ringbuffer, buf_in, size) == -1) { + // no space available + *errcode = MP_EAGAIN; + return MP_STREAM_ERROR; + } + *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); + switch (request) { + case MP_STREAM_POLL: { + mp_uint_t 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; + } + return ret; + } + case MP_STREAM_CLOSE: + case MP_STREAM_FLUSH: + return 0; + } + *errcode = MP_EINVAL; + return MP_STREAM_ERROR; +} + +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(µpython_ringbuffer_any_obj) }, + { MP_ROM_QSTR(MP_QSTR_reset), MP_ROM_PTR(µpython_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, µpython_ringbuffer_locals_dict + ); + +#endif diff --git a/py/ringbuf.h b/py/ringbuf.h index c8508c07edfa7..86a7912abbb21 100644 --- a/py/ringbuf.h +++ b/py/ringbuf.h @@ -28,9 +28,10 @@ #include #include +#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 new file mode 100644 index 0000000000000..162fe52d6ba64 --- /dev/null +++ b/tests/micropython/ringbuffer.py @@ -0,0 +1,38 @@ +# check that micropython.ringbuffer works correctly. + +import micropython + +try: + micropython.ringbuffer +except AttributeError: + print("SKIP") + raise SystemExit + +rb = micropython.ringbuffer(16) +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()) + +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 new file mode 100644 index 0000000000000..c2335f5a1684f --- /dev/null +++ b/tests/micropython/ringbuffer.py.exp @@ -0,0 +1,11 @@ + +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 new file mode 100644 index 0000000000000..8d63641a0b2d6 --- /dev/null +++ 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) +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 new file mode 100644 index 0000000000000..dfeb71d890fe0 --- /dev/null +++ b/tests/micropython/ringbuffer_async.py.exp @@ -0,0 +1,4 @@ +w 120 +b'ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123' +r 120 +True