shared-bindings/socket: add socket_recv_into

This commit is contained in:
Joshua Coats 2019-03-23 10:49:43 -07:00
parent fce63b17c6
commit e23bad3a3a

View File

@ -240,6 +240,52 @@ STATIC mp_obj_t socket_send(mp_obj_t self_in, mp_obj_t buf_in) {
}
STATIC MP_DEFINE_CONST_FUN_OBJ_2(socket_send_obj, socket_send);
// helper function for socket_recv and socket_recv_into to handle common operations of both
STATIC mp_int_t _socket_recv_into(mod_network_socket_obj_t *sock, byte *buf, mp_int_t len) {
int _errno;
mp_int_t ret = sock->nic_type->recv(sock, buf, len, &_errno);
if (ret == -1) {
mp_raise_OSError(_errno);
}
return len;
}
//| .. method:: recv_into(buffer[, bufsize])
//|
//| Reads some bytes from the connected remote address, writing
//| into the provided buffer. If bufsize <= len(buffer) is given,
//| a maximum of bufsize bytes will be read into the buffer. If no
//| valid value is given for bufsize, the default is the length of
//| the given buffer.
//|
//| Suits sockets of type SOCK_STREAM
//| Returns an int of number of bytes read.
//|
//| :param bytearray buffer: buffer to receive into
//| :param int bufsize: optionally, a maximum number of bytes to read.
STATIC mp_obj_t socket_recv_into(size_t n_args, const mp_obj_t *args) {
mod_network_socket_obj_t *self = MP_OBJ_TO_PTR(args[0]);
if (self->nic == MP_OBJ_NULL) {
// not connected
mp_raise_OSError(MP_ENOTCONN);
}
mp_buffer_info_t bufinfo;
mp_get_buffer_raise(args[1], &bufinfo, MP_BUFFER_WRITE);
mp_int_t len;
if (n_args == 3) {
len = mp_obj_get_int(args[2]);
}
if (n_args == 2 || (size_t) len > bufinfo.len) {
len = bufinfo.len;
}
mp_int_t ret = _socket_recv_into(self, (byte*)bufinfo.buf, len);
return mp_obj_new_int_from_uint(ret);
}
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(socket_recv_into_obj, 2, 3, socket_recv_into);
//| .. method:: recv(bufsize)
//|
//| Reads some bytes from the connected remote address.
@ -257,11 +303,7 @@ STATIC mp_obj_t socket_recv(mp_obj_t self_in, mp_obj_t len_in) {
mp_int_t len = mp_obj_get_int(len_in);
vstr_t vstr;
vstr_init_len(&vstr, len);
int _errno;
mp_int_t ret = self->nic_type->recv(self, (byte*)vstr.buf, len, &_errno);
if (ret == -1) {
mp_raise_OSError(_errno);
}
mp_int_t ret = _socket_recv_into(self, (byte*)vstr.buf, len);
if (ret == 0) {
return mp_const_empty_bytes;
}
@ -436,6 +478,7 @@ STATIC const mp_rom_map_elem_t socket_locals_dict_table[] = {
{ MP_ROM_QSTR(MP_QSTR_recv), MP_ROM_PTR(&socket_recv_obj) },
{ MP_ROM_QSTR(MP_QSTR_sendto), MP_ROM_PTR(&socket_sendto_obj) },
{ MP_ROM_QSTR(MP_QSTR_recvfrom), MP_ROM_PTR(&socket_recvfrom_obj) },
{ MP_ROM_QSTR(MP_QSTR_recv_into), MP_ROM_PTR(&socket_recv_into_obj) },
{ MP_ROM_QSTR(MP_QSTR_setsockopt), MP_ROM_PTR(&socket_setsockopt_obj) },
{ MP_ROM_QSTR(MP_QSTR_settimeout), MP_ROM_PTR(&socket_settimeout_obj) },
{ MP_ROM_QSTR(MP_QSTR_setblocking), MP_ROM_PTR(&socket_setblocking_obj) },