py/objstringio: Implement MP_STREAM_SEEK ioctl and add seek() method.

1 files changed, 24 insertions, 3 deletions

diff --git a/py/objstringio.c b/py/objstringio.c
index a75a334332..eb2e516bb3 100644
--- a/py/objstringio.c
+++ b/py/objstringio.c
@@ -75,13 +75,18 @@ 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_uint_t remaining = o->vstr->alloc - o->pos;
-    if (size > remaining) {
+    mp_int_t remaining = o->vstr->alloc - o->pos;
+    mp_uint_t org_len = o->vstr->len;
+    if ((mp_int_t)size > remaining) {
         // Take all what's already allocated...
         o->vstr->len = o->vstr->alloc;
         // ... and add more
         vstr_add_len(o->vstr, size - remaining);
     }
+    // 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) {
@@ -92,8 +97,23 @@ STATIC mp_uint_t stringio_write(mp_obj_t o_in, const void *buf, mp_uint_t size,
 
 STATIC mp_uint_t stringio_ioctl(mp_obj_t o_in, mp_uint_t request, uintptr_t arg, int *errcode) {
     (void)errcode;
-//    mp_obj_stringio_t *o = MP_OBJ_TO_PTR(o_in);
+    mp_obj_stringio_t *o = MP_OBJ_TO_PTR(o_in);
     switch (request) {
+        case MP_STREAM_SEEK: {
+            struct mp_stream_seek_t *s = (struct mp_stream_seek_t*)arg;
+            mp_uint_t ref = 0;
+            switch (s->whence) {
+                case 1: // SEEK_CUR
+                    ref = o->pos;
+                    break;
+                case 2: // SEEK_END
+                    ref = o->vstr->len;
+                    break;
+            }
+            o->pos = ref + s->offset;
+            s->offset = o->pos;
+            return 0;
+        }
         case MP_STREAM_FLUSH:
             return 0;
         default:
@@ -160,6 +180,7 @@ STATIC const mp_rom_map_elem_t stringio_locals_dict_table[] = {
     { MP_ROM_QSTR(MP_QSTR_readall), MP_ROM_PTR(&mp_stream_readall_obj) },
     { MP_ROM_QSTR(MP_QSTR_readline), MP_ROM_PTR(&mp_stream_unbuffered_readline_obj) },
     { MP_ROM_QSTR(MP_QSTR_write), MP_ROM_PTR(&mp_stream_write_obj) },
+    { MP_ROM_QSTR(MP_QSTR_seek), MP_ROM_PTR(&mp_stream_seek_obj) },
     { MP_ROM_QSTR(MP_QSTR_flush), MP_ROM_PTR(&mp_stream_flush_obj) },
     { MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&stringio_close_obj) },
     { MP_ROM_QSTR(MP_QSTR_getvalue), MP_ROM_PTR(&stringio_getvalue_obj) },