From 50d3a9df671e5154c13a8218ad0e90fb9e78171f Mon Sep 17 00:00:00 2001 From: Paul Sokolovsky Date: Thu, 2 Feb 2017 00:33:43 +0300 Subject: [PATCH] py/objstringio: Allow to specify initial capacity by passing numeric argument. E.g. uio.BytesIO(100) will allocate buffer with 100 bytes of space. --- py/objstringio.c | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/py/objstringio.c b/py/objstringio.c index a430fca3b7..a77ffae246 100644 --- a/py/objstringio.c +++ b/py/objstringio.c @@ -147,21 +147,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;