diff --git a/py/objstringio.c b/py/objstringio.c index 9f4adeebb..645c441cb 100644 --- a/py/objstringio.c +++ b/py/objstringio.c @@ -68,10 +68,23 @@ STATIC mp_uint_t stringio_read(mp_obj_t o_in, void *buf, mp_uint_t size, int *er return size; } +STATIC void stringio_copy_on_write(mp_obj_stringio_t *o) { + const void *buf = o->vstr->buf; + o->vstr->buf = m_new(char, o->vstr->len); + memcpy(o->vstr->buf, buf, o->vstr->len); + o->vstr->fixed_buf = false; + o->ref_obj = MP_OBJ_NULL; +} + STATIC mp_uint_t stringio_write(mp_obj_t o_in, const void *buf, mp_uint_t size, int *errcode) { (void)errcode; mp_obj_stringio_t *o = MP_OBJ_TO_PTR(o_in); check_stringio_is_open(o); + + if (o->vstr->fixed_buf) { + stringio_copy_on_write(o); + } + mp_uint_t new_pos = o->pos + size; if (new_pos < size) { // Writing bytes will overflow o->pos beyond limit of mp_uint_t. @@ -155,11 +168,11 @@ 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, mp_uint_t alloc) { +STATIC mp_obj_stringio_t *stringio_new(const mp_obj_type_t *type) { mp_obj_stringio_t *o = m_new_obj(mp_obj_stringio_t); o->base.type = type; - o->vstr = vstr_new(alloc); o->pos = 0; + o->ref_obj = MP_OBJ_NULL; return o; } @@ -170,17 +183,28 @@ STATIC mp_obj_t stringio_make_new(const mp_obj_type_t *type_in, size_t n_args, s bool initdata = false; mp_buffer_info_t bufinfo; + mp_obj_stringio_t *o = stringio_new(type_in); + if (n_args > 0) { 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); + + if (MP_OBJ_IS_STR_OR_BYTES(args[0])) { + o->vstr = m_new_obj(vstr_t); + vstr_init_fixed_buf(o->vstr, bufinfo.len, bufinfo.buf); + o->vstr->len = bufinfo.len; + o->ref_obj = args[0]; + return MP_OBJ_FROM_PTR(o); + } + sz = bufinfo.len; initdata = true; } } - mp_obj_stringio_t *o = stringio_new(type_in, sz); + o->vstr = vstr_new(sz); if (initdata) { stringio_write(MP_OBJ_FROM_PTR(o), bufinfo.buf, bufinfo.len, NULL); diff --git a/py/objstringio.h b/py/objstringio.h index 853bfb11b..56738f4e4 100644 --- a/py/objstringio.h +++ b/py/objstringio.h @@ -33,6 +33,8 @@ typedef struct _mp_obj_stringio_t { vstr_t *vstr; // StringIO has single pointer used for both reading and writing mp_uint_t pos; + // Underlying object buffered by this StringIO + mp_obj_t ref_obj; } mp_obj_stringio_t; #endif // MICROPY_INCLUDED_PY_OBJSTRINGIO_H diff --git a/tests/io/bytesio_cow.py b/tests/io/bytesio_cow.py new file mode 100644 index 000000000..92654a000 --- /dev/null +++ b/tests/io/bytesio_cow.py @@ -0,0 +1,20 @@ +# Make sure that write operations on io.BytesIO don't +# change original object it was constructed from. +try: + import uio as io +except ImportError: + import io + +b = b"foobar" + +a = io.BytesIO(b) +a.write(b"1") +print(b) +print(a.getvalue()) + +b = bytearray(b"foobar") + +a = io.BytesIO(b) +a.write(b"1") +print(b) +print(a.getvalue()) diff --git a/tests/micropython/heapalloc_bytesio2.py b/tests/micropython/heapalloc_bytesio2.py new file mode 100644 index 000000000..cd76f5807 --- /dev/null +++ b/tests/micropython/heapalloc_bytesio2.py @@ -0,0 +1,20 @@ +# Creating BytesIO from immutable object should not immediately +# copy its content. +try: + import uio + import micropython + micropython.mem_total +except (ImportError, AttributeError): + print("SKIP") + raise SystemExit + + +data = b"1234" * 256 + +before = micropython.mem_total() + +buf = uio.BytesIO(data) + +after = micropython.mem_total() + +print(after - before < len(data)) diff --git a/tests/micropython/heapalloc_bytesio2.py.exp b/tests/micropython/heapalloc_bytesio2.py.exp new file mode 100644 index 000000000..0ca95142b --- /dev/null +++ b/tests/micropython/heapalloc_bytesio2.py.exp @@ -0,0 +1 @@ +True