diff options
Diffstat (limited to 'py/objstringio.c')
| -rw-r--r-- | py/objstringio.c | 49 |
1 files changed, 35 insertions, 14 deletions
diff --git a/py/objstringio.c b/py/objstringio.c index a430fca3b7..9f4adeebbf 100644 --- a/py/objstringio.c +++ b/py/objstringio.c @@ -39,7 +39,7 @@ #if MICROPY_CPYTHON_COMPAT STATIC void check_stringio_is_open(const mp_obj_stringio_t *o) { if (o->vstr == NULL) { - mp_raise_msg(&mp_type_ValueError, "I/O operation on closed file"); + mp_raise_ValueError("I/O operation on closed file"); } } #else @@ -56,6 +56,9 @@ STATIC mp_uint_t stringio_read(mp_obj_t o_in, void *buf, mp_uint_t size, int *er (void)errcode; mp_obj_stringio_t *o = MP_OBJ_TO_PTR(o_in); check_stringio_is_open(o); + if (o->vstr->len <= o->pos) { // read to EOF, or seeked to EOF or beyond + return 0; + } mp_uint_t remaining = o->vstr->len - o->pos; if (size > remaining) { size = remaining; @@ -69,22 +72,27 @@ STATIC mp_uint_t stringio_write(mp_obj_t o_in, const void *buf, mp_uint_t size, (void)errcode; mp_obj_stringio_t *o = MP_OBJ_TO_PTR(o_in); check_stringio_is_open(o); - mp_int_t remaining = o->vstr->alloc - o->pos; + mp_uint_t new_pos = o->pos + size; + if (new_pos < size) { + // Writing <size> bytes will overflow o->pos beyond limit of mp_uint_t. + *errcode = MP_EFBIG; + return MP_STREAM_ERROR; + } mp_uint_t org_len = o->vstr->len; - if ((mp_int_t)size > remaining) { + if (new_pos > o->vstr->alloc) { // Take all what's already allocated... o->vstr->len = o->vstr->alloc; // ... and add more - vstr_add_len(o->vstr, size - remaining); + vstr_add_len(o->vstr, new_pos - o->vstr->alloc); } // If there was a seek past EOF, clear the hole if (o->pos > org_len) { memset(o->vstr->buf + org_len, 0, o->pos - org_len); } memcpy(o->vstr->buf + o->pos, buf, size); - o->pos += size; - if (o->pos > o->vstr->len) { - o->vstr->len = o->pos; + o->pos = new_pos; + if (new_pos > o->vstr->len) { + o->vstr->len = new_pos; } return size; } @@ -147,21 +155,34 @@ STATIC mp_obj_t stringio___exit__(size_t n_args, const mp_obj_t *args) { } STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(stringio___exit___obj, 4, 4, stringio___exit__); -STATIC mp_obj_stringio_t *stringio_new(const mp_obj_type_t *type) { +STATIC mp_obj_stringio_t *stringio_new(const mp_obj_type_t *type, mp_uint_t alloc) { mp_obj_stringio_t *o = m_new_obj(mp_obj_stringio_t); o->base.type = type; - o->vstr = vstr_new(16); + o->vstr = vstr_new(alloc); o->pos = 0; return o; } STATIC mp_obj_t stringio_make_new(const mp_obj_type_t *type_in, size_t n_args, size_t n_kw, const mp_obj_t *args) { (void)n_kw; // TODO check n_kw==0 - mp_obj_stringio_t *o = stringio_new(type_in); + + mp_uint_t sz = 16; + bool initdata = false; + mp_buffer_info_t bufinfo; if (n_args > 0) { - mp_buffer_info_t bufinfo; - mp_get_buffer_raise(args[0], &bufinfo, MP_BUFFER_READ); + if (MP_OBJ_IS_INT(args[0])) { + sz = mp_obj_get_int(args[0]); + } else { + mp_get_buffer_raise(args[0], &bufinfo, MP_BUFFER_READ); + sz = bufinfo.len; + initdata = true; + } + } + + mp_obj_stringio_t *o = stringio_new(type_in, sz); + + if (initdata) { stringio_write(MP_OBJ_FROM_PTR(o), bufinfo.buf, bufinfo.len, NULL); // Cur ptr is always at the beginning of buffer at the construction o->pos = 0; @@ -202,7 +223,7 @@ const mp_obj_type_t mp_type_stringio = { .name = MP_QSTR_StringIO, .print = stringio_print, .make_new = stringio_make_new, - .getiter = mp_identity, + .getiter = mp_identity_getiter, .iternext = mp_stream_unbuffered_iter, .protocol = &stringio_stream_p, .locals_dict = (mp_obj_dict_t*)&stringio_locals_dict, @@ -214,7 +235,7 @@ const mp_obj_type_t mp_type_bytesio = { .name = MP_QSTR_BytesIO, .print = stringio_print, .make_new = stringio_make_new, - .getiter = mp_identity, + .getiter = mp_identity_getiter, .iternext = mp_stream_unbuffered_iter, .protocol = &bytesio_stream_p, .locals_dict = (mp_obj_dict_t*)&stringio_locals_dict, |