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 <damien@micropython.org>
This commit is contained in:
Damien George
2025-09-13 19:33:34 +10:00
parent 10b3d96b60
commit 2373340aa3

View File

@@ -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);