diff options
| -rw-r--r-- | py/objstringio.c | 23 |
1 files 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; |