py/objstringio: Allow to specify initial capacity by passing numeric argument.

E.g. uio.BytesIO(100) will allocate buffer with 100 bytes of space.
This commit is contained in:
Paul Sokolovsky 2017-02-02 00:33:43 +03:00
parent aee74a1dae
commit 50d3a9df67
1 changed files with 18 additions and 5 deletions

View File

@ -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;