From 2373340aa342dbd30e1482cf431895f808d863b2 Mon Sep 17 00:00:00 2001 From: Damien George Date: Sat, 13 Sep 2025 19:33:34 +1000 Subject: [PATCH] py/stream: Reuse write implementation for readinto. This commit refactors some common code in the core stream implementation, to reduce code size while retaining the same functionality. With the factoring, `readinto`/`readinto1` could now support an additional 4th argument (like write) but it's best not to introduce even more CPython incompatibility, so they are left as having a maximum of 3 args. Signed-off-by: Damien George --- py/stream.c | 46 +++++++++++++--------------------------------- 1 file changed, 13 insertions(+), 33 deletions(-) diff --git a/py/stream.c b/py/stream.c index 008c968452..93bdac42c8 100644 --- a/py/stream.c +++ b/py/stream.c @@ -261,9 +261,14 @@ void mp_stream_write_adaptor(void *self, const char *buf, size_t len) { mp_stream_write(MP_OBJ_FROM_PTR(self), buf, len, MP_STREAM_RW_WRITE); } -static mp_obj_t stream_write_generic(size_t n_args, const mp_obj_t *args, byte flags) { +static mp_obj_t stream_readinto_write_generic(size_t n_args, const mp_obj_t *args, byte flags) { mp_buffer_info_t bufinfo; - mp_get_buffer_raise(args[1], &bufinfo, MP_BUFFER_READ); + mp_get_buffer_raise(args[1], &bufinfo, (flags & MP_STREAM_RW_WRITE) ? MP_BUFFER_READ : MP_BUFFER_WRITE); + + // CPython extension, allow optional maximum length and offset: + // - stream.operation(buf, max_len) + // - stream.operation(buf, off, max_len) + // Similar to https://docs.python.org/3/library/socket.html#socket.socket.recv_into size_t max_len = (size_t)-1; size_t off = 0; if (n_args == 3) { @@ -276,53 +281,28 @@ static mp_obj_t stream_write_generic(size_t n_args, const mp_obj_t *args, byte f } } bufinfo.len -= off; + + // Perform the readinto or write operation. return mp_stream_write(args[0], (byte *)bufinfo.buf + off, MIN(bufinfo.len, max_len), flags); } static mp_obj_t stream_write_method(size_t n_args, const mp_obj_t *args) { - return stream_write_generic(n_args, args, MP_STREAM_RW_WRITE); + return stream_readinto_write_generic(n_args, args, MP_STREAM_RW_WRITE); } MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_write_obj, 2, 4, stream_write_method); static mp_obj_t stream_write1_method(size_t n_args, const mp_obj_t *args) { - return stream_write_generic(n_args, args, MP_STREAM_RW_WRITE | MP_STREAM_RW_ONCE); + return stream_readinto_write_generic(n_args, args, MP_STREAM_RW_WRITE | MP_STREAM_RW_ONCE); } MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_write1_obj, 2, 4, stream_write1_method); -static mp_obj_t stream_readinto_generic(size_t n_args, const mp_obj_t *args, byte flags) { - mp_buffer_info_t bufinfo; - mp_get_buffer_raise(args[1], &bufinfo, MP_BUFFER_WRITE); - - // CPython extension: if 2nd arg is provided, that's max len to read, - // instead of full buffer. Similar to - // https://docs.python.org/3/library/socket.html#socket.socket.recv_into - mp_uint_t len = bufinfo.len; - if (n_args > 2) { - len = mp_obj_get_int(args[2]); - if (len > bufinfo.len) { - len = bufinfo.len; - } - } - - int error; - mp_uint_t out_sz = mp_stream_rw(args[0], bufinfo.buf, len, &error, flags); - if (error != 0) { - if (mp_is_nonblocking_error(error)) { - return mp_const_none; - } - mp_raise_OSError(error); - } else { - return MP_OBJ_NEW_SMALL_INT(out_sz); - } -} - static mp_obj_t stream_readinto(size_t n_args, const mp_obj_t *args) { - return stream_readinto_generic(n_args, args, MP_STREAM_RW_READ); + return stream_readinto_write_generic(n_args, args, MP_STREAM_RW_READ); } MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_readinto_obj, 2, 3, stream_readinto); static mp_obj_t stream_readinto1(size_t n_args, const mp_obj_t *args) { - return stream_readinto_generic(n_args, args, MP_STREAM_RW_READ | MP_STREAM_RW_ONCE); + return stream_readinto_write_generic(n_args, args, MP_STREAM_RW_READ | MP_STREAM_RW_ONCE); } MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_readinto1_obj, 2, 3, stream_readinto1);