diff --git a/docs/library/micropython.rst b/docs/library/micropython.rst index b17dfa9a75a48..11953e434f8ca 100644 --- a/docs/library/micropython.rst +++ b/docs/library/micropython.rst @@ -112,6 +112,25 @@ Functions This function can be used to prevent the capturing of Ctrl-C on the incoming stream of characters that is usually used for the REPL, in case that stream is used for other purposes. + +.. function:: RingIO(size) +.. function:: RingIO(buffer) + :noindex: + + Provides a fixed-size ringbuffer for bytes with a stream interface. Can be + considered like a fifo queue variant of `io.BytesIO`. + + Can be created with integer size provided and a suitable buffer will be allocated. + Alternatively a `bytearray`, `memoryview` or similar object can be provided at + creation for in-place use. + + The classic ringbuffer algorithm is used which allows for any size buffer to be + used however one byte will be consumed for tracking. If initialised with an integer + size this will be accounted for, for example ``RingIO(16)`` will allocate a 17 byte + buffer internally so it can hold 16 bytes of data. + When passing in a pre-allocated buffer however one byte less than its original + length will be available for storage, eg. ``RingIO(bytearray(16))`` will only + hold 15 bytes of data. .. function:: schedule(func, arg) diff --git a/py/modmicropython.c b/py/modmicropython.c index af6ad01795f67..c101e128c60d8 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 @@ -200,6 +201,9 @@ static const mp_rom_map_elem_t mp_module_micropython_globals_table[] = { #if MICROPY_KBD_EXCEPTION { MP_ROM_QSTR(MP_QSTR_kbd_intr), MP_ROM_PTR(&mp_micropython_kbd_intr_obj) }, #endif + #if MICROPY_PY_MICROPYTHON_RINGIO + { MP_ROM_QSTR(MP_QSTR_RingIO), MP_ROM_PTR(&mp_type_micropython_ringio) }, + #endif #if MICROPY_ENABLE_SCHEDULER { MP_ROM_QSTR(MP_QSTR_schedule), MP_ROM_PTR(&mp_micropython_schedule_obj) }, #endif diff --git a/py/mpconfig.h b/py/mpconfig.h index 94fdca7d7a97d..21454abdef59b 100644 --- a/py/mpconfig.h +++ b/py/mpconfig.h @@ -1299,6 +1299,11 @@ typedef double mp_float_t; #define MICROPY_PY_MICROPYTHON_HEAP_LOCKED (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EVERYTHING) #endif +// Support for micropython.RingIO() +#ifndef MICROPY_PY_MICROPYTHON_RINGIO +#define MICROPY_PY_MICROPYTHON_RINGIO (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES) +#endif + // Whether to provide "array" module. Note that large chunk of the // underlying code is shared with "bytearray" builtin type, so to // get real savings, it should be disabled too. diff --git a/py/ringbuf.c b/py/ringbuf.c index 10dca62081e8c..5914ca6fab1eb 100644 --- a/py/ringbuf.c +++ b/py/ringbuf.c @@ -75,14 +75,8 @@ int ringbuf_put16(ringbuf_t *r, uint16_t v) { return 0; } -// Returns: -// 0: Success -// -1: Not enough data available to complete read (try again later) -// -2: Requested read is larger than buffer - will never succeed -int ringbuf_get_bytes(ringbuf_t *r, uint8_t *data, size_t data_len) { - if (ringbuf_avail(r) < data_len) { - return (r->size <= data_len) ? -2 : -1; - } +static inline MP_ALWAYSINLINE void _ringbuf_memcpy_get(ringbuf_t *r, uint8_t *data, size_t data_len) { + // Ensure bounds / space checking is done before running this uint32_t iget = r->iget; uint32_t iget_a = (iget + data_len) % r->size; uint8_t *datap = data; @@ -94,17 +88,10 @@ int ringbuf_get_bytes(ringbuf_t *r, uint8_t *data, size_t data_len) { } memcpy(datap, r->buf + iget, iget_a - iget); r->iget = iget_a; - return 0; } -// Returns: -// 0: Success -// -1: Not enough free space available to complete write (try again later) -// -2: Requested write is larger than buffer - will never succeed -int ringbuf_put_bytes(ringbuf_t *r, const uint8_t *data, size_t data_len) { - if (ringbuf_free(r) < data_len) { - return (r->size <= data_len) ? -2 : -1; - } +static inline MP_ALWAYSINLINE void _ringbuf_memcpy_put(ringbuf_t *r, const uint8_t *data, size_t data_len) { + // Ensure bounds / space checking is done before running this uint32_t iput = r->iput; uint32_t iput_a = (iput + data_len) % r->size; const uint8_t *datap = data; @@ -116,5 +103,130 @@ int ringbuf_put_bytes(ringbuf_t *r, const uint8_t *data, size_t data_len) { } memcpy(r->buf + iput, datap, iput_a - iput); r->iput = iput_a; +} + +// Returns: +// 0: Success +// -1: Not enough data available to complete read (try again later) +// -2: Requested read is larger than buffer - will never succeed +int ringbuf_get_bytes(ringbuf_t *r, uint8_t *data, size_t data_len) { + if (ringbuf_avail(r) < data_len) { + return (r->size <= data_len) ? -2 : -1; + } + _ringbuf_memcpy_get(r, data, data_len); return 0; } + +// Returns: +// 0: Success +// -1: Not enough free space available to complete write (try again later) +// -2: Requested write is larger than buffer - will never succeed +int ringbuf_put_bytes(ringbuf_t *r, const uint8_t *data, size_t data_len) { + if (ringbuf_free(r) < data_len) { + return (r->size <= data_len) ? -2 : -1; + } + _ringbuf_memcpy_put(r, data, data_len); + return 0; +} + +#if MICROPY_PY_MICROPYTHON_RINGIO +#include "py/runtime.h" +#include "py/stream.h" +#include "py/mphal.h" + +typedef struct _micropython_ringio_obj_t { + mp_obj_base_t base; + ringbuf_t ringbuffer; +} micropython_ringio_obj_t; + +static mp_obj_t micropython_ringio_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_ringio_obj_t *self = mp_obj_malloc(micropython_ringio_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 { + // 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_ringio_read(mp_obj_t self_in, void *buf_in, mp_uint_t size, int *errcode) { + micropython_ringio_obj_t *self = MP_OBJ_TO_PTR(self_in); + size = MIN(size, ringbuf_avail(&self->ringbuffer)); + _ringbuf_memcpy_get(&(self->ringbuffer), buf_in, size); + *errcode = 0; + return size; +} + +static mp_uint_t micropython_ringio_write(mp_obj_t self_in, const void *buf_in, mp_uint_t size, int *errcode) { + micropython_ringio_obj_t *self = MP_OBJ_TO_PTR(self_in); + size = MIN(size, ringbuf_free(&self->ringbuffer)); + _ringbuf_memcpy_put(&(self->ringbuffer), buf_in, size); + *errcode = 0; + return size; +} + +static mp_uint_t micropython_ringio_ioctl(mp_obj_t self_in, mp_uint_t request, uintptr_t arg, int *errcode) { + micropython_ringio_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: + return 0; + } + *errcode = MP_EINVAL; + return MP_STREAM_ERROR; +} + +static mp_obj_t micropython_ringio_any(mp_obj_t self_in) { + micropython_ringio_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_ringio_any_obj, micropython_ringio_any); + +static const mp_rom_map_elem_t micropython_ringio_locals_dict_table[] = { + { MP_ROM_QSTR(MP_QSTR_any), MP_ROM_PTR(µpython_ringio_any_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_ringio_locals_dict, micropython_ringio_locals_dict_table); + +static const mp_stream_p_t ringio_stream_p = { + .read = micropython_ringio_read, + .write = micropython_ringio_write, + .ioctl = micropython_ringio_ioctl, + .is_text = false, +}; + +MP_DEFINE_CONST_OBJ_TYPE( + mp_type_micropython_ringio, + MP_QSTR_RingIO, + MP_TYPE_FLAG_NONE, + make_new, micropython_ringio_make_new, + protocol, &ringio_stream_p, + locals_dict, µpython_ringio_locals_dict + ); + +#endif // MICROPY_PY_MICROPYTHON_RINGIO diff --git a/py/ringbuf.h b/py/ringbuf.h index c8508c07edfa7..aba42dfc496de 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_RINGIO +#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_RINGIO +extern const mp_obj_type_t mp_type_micropython_ringio; +#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..caadd08dba299 --- /dev/null +++ b/tests/micropython/ringbuffer.py @@ -0,0 +1,48 @@ +# check that micropython.RingIO works correctly. + +import micropython + +try: + micropython.RingIO +except AttributeError: + print("SKIP") + raise SystemExit + +rb = micropython.RingIO(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)) + +# try to write more data than can fit at one go +print(rb.write(b"\x00\x01" * 10)) +print(rb.write(b"\x00")) +print(rb.read()) + + +ba = bytearray(17) +rb = micropython.RingIO(ba) +print(rb) +print(rb.write(b"\x00\x01" * 10)) +print(rb.write(b"\x00")) +print(rb.read()) + +try: + # size must be int. + micropython.RingIO(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..b27ea472c7d89 --- /dev/null +++ b/tests/micropython/ringbuffer.py.exp @@ -0,0 +1,16 @@ + +0 +1 +2 +b'\x00\x00' +0 +b'\x00\x01' +b'' +16 +0 +b'\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01' + +16 +0 +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..78f74c40bef53 --- /dev/null +++ b/tests/micropython/ringbuffer_async.py @@ -0,0 +1,37 @@ +# check that micropython.RingIO works correctly with asyncio.Stream. + +import micropython + +try: + import asyncio + + asyncio.StreamWriter + micropython.RingIO +except (AttributeError, ImportError): + print("SKIP") + raise SystemExit + +rb = micropython.RingIO(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