mirror of
https://github.com/micropython/micropython.git
synced 2025-12-15 17:30:14 +01:00
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:
46
py/stream.c
46
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);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user