diff --git a/py/lexer.c b/py/lexer.c index f7f9c631f3..da8967b163 100644 --- a/py/lexer.c +++ b/py/lexer.c @@ -614,7 +614,7 @@ mp_lexer_t *mp_lexer_new(const char *src_name, void *stream_data, mp_lexer_strea lex->num_indent_level = 1; lex->indent_level = m_new(uint16_t, lex->alloc_indent_level); lex->indent_level[0] = 0; - vstr_init(&lex->vstr); + vstr_init(&lex->vstr, 32); // preload characters lex->chr0 = stream_next_char(stream_data); diff --git a/py/misc.h b/py/misc.h index 22158c71a7..e985bc5253 100644 --- a/py/misc.h +++ b/py/misc.h @@ -59,15 +59,19 @@ typedef struct _vstr_t { bool had_error; } vstr_t; -void vstr_init(vstr_t *vstr); +void vstr_init(vstr_t *vstr, int alloc); void vstr_clear(vstr_t *vstr); vstr_t *vstr_new(void); +vstr_t *vstr_new_size(int alloc); void vstr_free(vstr_t *vstr); void vstr_reset(vstr_t *vstr); bool vstr_had_error(vstr_t *vstr); char *vstr_str(vstr_t *vstr); int vstr_len(vstr_t *vstr); void vstr_hint_size(vstr_t *vstr, int size); +char *vstr_extend(vstr_t *vstr, int size); +bool vstr_set_size(vstr_t *vstr, int size); +bool vstr_shrink(vstr_t *vstr); char *vstr_add_len(vstr_t *vstr, int len); void vstr_add_byte(vstr_t *vstr, byte v); void vstr_add_char(vstr_t *vstr, unichar chr); diff --git a/py/stream.c b/py/stream.c index f883efcd25..48cd0bcd6c 100644 --- a/py/stream.c +++ b/py/stream.c @@ -51,5 +51,47 @@ static mp_obj_t stream_write(mp_obj_t self_in, mp_obj_t arg) { } } +// TODO: should be in mpconfig.h +#define READ_SIZE 256 +static mp_obj_t stream_readall(mp_obj_t self_in) { + struct _mp_obj_base_t *o = (struct _mp_obj_base_t *)self_in; + if (o->type->stream_p.read == NULL) { + // CPython: io.UnsupportedOperation, OSError subclass + nlr_jump(mp_obj_new_exception_msg(MP_QSTR_OSError, "Operation not supported")); + } + + int total_size = 0; + vstr_t *vstr = vstr_new_size(READ_SIZE); + char *buf = vstr_str(vstr); + char *p = buf; + int error; + int current_read = READ_SIZE; + while (true) { + machine_int_t out_sz = o->type->stream_p.read(self_in, p, current_read, &error); + if (out_sz == -1) { + nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_OSError, "[Errno %d]", error)); + } + if (out_sz == 0) { + break; + } + total_size += out_sz; + if (out_sz < current_read) { + current_read -= out_sz; + p += out_sz; + } else { + current_read = READ_SIZE; + p = vstr_extend(vstr, current_read); + if (p == NULL) { + // TODO + nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_OSError/*MP_QSTR_RuntimeError*/, "Out of memory")); + } + } + } + vstr_set_size(vstr, total_size + 1); // TODO: for \0 + buf[total_size] = 0; + return mp_obj_new_str(qstr_from_str_take(buf, total_size + 1)); +} + MP_DEFINE_CONST_FUN_OBJ_2(mp_stream_read_obj, stream_read); +MP_DEFINE_CONST_FUN_OBJ_1(mp_stream_readall_obj, stream_readall); MP_DEFINE_CONST_FUN_OBJ_2(mp_stream_write_obj, stream_write); diff --git a/py/stream.h b/py/stream.h index 8a579b7621..b1c3c72786 100644 --- a/py/stream.h +++ b/py/stream.h @@ -1,2 +1,3 @@ extern const mp_obj_fun_native_t mp_stream_read_obj; +extern const mp_obj_fun_native_t mp_stream_readall_obj; extern const mp_obj_fun_native_t mp_stream_write_obj; diff --git a/py/vstr.c b/py/vstr.c index 80841b24ca..18e6d2d0a8 100644 --- a/py/vstr.c +++ b/py/vstr.c @@ -6,8 +6,8 @@ // returned value is always at least 1 greater than argument #define ROUND_ALLOC(a) (((a) & ((~0) - 7)) + 8) -void vstr_init(vstr_t *vstr) { - vstr->alloc = 32; +void vstr_init(vstr_t *vstr, int alloc) { + vstr->alloc = alloc; vstr->len = 0; vstr->buf = m_new(char, vstr->alloc); if (vstr->buf == NULL) { @@ -28,7 +28,16 @@ vstr_t *vstr_new(void) { if (vstr == NULL) { return NULL; } - vstr_init(vstr); + vstr_init(vstr, 32); + return vstr; +} + +vstr_t *vstr_new_size(int alloc) { + vstr_t *vstr = m_new(vstr_t, 1); + if (vstr == NULL) { + return NULL; + } + vstr_init(vstr, alloc); return vstr; } @@ -63,6 +72,36 @@ int vstr_len(vstr_t *vstr) { return vstr->len; } +// Extend vstr strictly to by requested size, return pointer to newly added chunk +char *vstr_extend(vstr_t *vstr, int size) { + char *new_buf = m_renew(char, vstr->buf, vstr->alloc, vstr->alloc + size); + if (new_buf == NULL) { + vstr->had_error = true; + return NULL; + } + char *p = new_buf + vstr->alloc; + vstr->alloc += size; + vstr->buf = new_buf; + return p; +} + +// Shrink vstr to be given size +bool vstr_set_size(vstr_t *vstr, int size) { + char *new_buf = m_renew(char, vstr->buf, vstr->alloc, size); + if (new_buf == NULL) { + vstr->had_error = true; + return false; + } + vstr->buf = new_buf; + vstr->alloc = vstr->len = size; + return true; +} + +// Shrink vstr allocation to its actual length +bool vstr_shrink(vstr_t *vstr) { + return vstr_set_size(vstr, vstr->len); +} + bool vstr_ensure_extra(vstr_t *vstr, int size) { if (vstr->len + size + 1 > vstr->alloc) { int new_alloc = ROUND_ALLOC((vstr->len + size + 1) * 2); diff --git a/stm/main.c b/stm/main.c index 2085182c29..4d49b8d3d2 100644 --- a/stm/main.c +++ b/stm/main.c @@ -390,7 +390,7 @@ void do_repl(void) { stdout_tx_str("Type \"help()\" for more information.\r\n"); vstr_t line; - vstr_init(&line); + vstr_init(&line, 32); for (;;) { vstr_reset(&line); diff --git a/unix/file.c b/unix/file.c index 398aac0759..204dc1b7ef 100644 --- a/unix/file.c +++ b/unix/file.c @@ -90,6 +90,7 @@ static mp_obj_t fdfile_make_new(mp_obj_t type_in, int n_args, const mp_obj_t *ar static const mp_method_t rawfile_type_methods[] = { { "read", &mp_stream_read_obj }, + { "readall", &mp_stream_readall_obj }, { "write", &mp_stream_write_obj }, { "close", &fdfile_close_obj }, { NULL, NULL },