Merge pull request #173 from pfalcon/file-readall

Generic implementation if stream readall() method, immediately reused in unix io.FileIO implementation

5 files changed, 91 insertions, 5 deletions

diff --git a/py/lexer.c b/py/lexer.c
index f7f9c631f3..da8967b163 100644
--- a/py/lexer.c
+++ b/py/lexer.c
@@ -614,7 +614,7 @@ mp_lexer_t *mp_lexer_new(const char *src_name, void *stream_data, mp_lexer_strea
     lex->num_indent_level = 1;
     lex->indent_level = m_new(uint16_t, lex->alloc_indent_level);
     lex->indent_level[0] = 0;
-    vstr_init(&lex->vstr);
+    vstr_init(&lex->vstr, 32);
 
     // preload characters
     lex->chr0 = stream_next_char(stream_data);
diff --git a/py/misc.h b/py/misc.h
index 22158c71a7..e985bc5253 100644
--- a/py/misc.h
+++ b/py/misc.h
@@ -59,15 +59,19 @@ typedef struct _vstr_t {
     bool had_error;
 } vstr_t;
 
-void vstr_init(vstr_t *vstr);
+void vstr_init(vstr_t *vstr, int alloc);
 void vstr_clear(vstr_t *vstr);
 vstr_t *vstr_new(void);
+vstr_t *vstr_new_size(int alloc);
 void vstr_free(vstr_t *vstr);
 void vstr_reset(vstr_t *vstr);
 bool vstr_had_error(vstr_t *vstr);
 char *vstr_str(vstr_t *vstr);
 int vstr_len(vstr_t *vstr);
 void vstr_hint_size(vstr_t *vstr, int size);
+char  *vstr_extend(vstr_t *vstr, int size);
+bool vstr_set_size(vstr_t *vstr, int size);
+bool vstr_shrink(vstr_t *vstr);
 char *vstr_add_len(vstr_t *vstr, int len);
 void vstr_add_byte(vstr_t *vstr, byte v);
 void vstr_add_char(vstr_t *vstr, unichar chr);
diff --git a/py/stream.c b/py/stream.c
index f883efcd25..48cd0bcd6c 100644
--- a/py/stream.c
+++ b/py/stream.c
@@ -51,5 +51,47 @@ static mp_obj_t stream_write(mp_obj_t self_in, mp_obj_t arg) {
     }
 }
 
+// TODO: should be in mpconfig.h
+#define READ_SIZE 256
+static mp_obj_t stream_readall(mp_obj_t self_in) {
+    struct _mp_obj_base_t *o = (struct _mp_obj_base_t *)self_in;
+    if (o->type->stream_p.read == NULL) {
+        // CPython: io.UnsupportedOperation, OSError subclass
+        nlr_jump(mp_obj_new_exception_msg(MP_QSTR_OSError, "Operation not supported"));
+    }
+
+    int total_size = 0;
+    vstr_t *vstr = vstr_new_size(READ_SIZE);
+    char *buf = vstr_str(vstr);
+    char *p = buf;
+    int error;
+    int current_read = READ_SIZE;
+    while (true) {
+        machine_int_t out_sz = o->type->stream_p.read(self_in, p, current_read, &error);
+        if (out_sz == -1) {
+            nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_OSError, "[Errno %d]", error));
+        }
+        if (out_sz == 0) {
+            break;
+        }
+        total_size += out_sz;
+        if (out_sz < current_read) {
+            current_read -= out_sz;
+            p += out_sz;
+        } else {
+            current_read = READ_SIZE;
+            p = vstr_extend(vstr, current_read);
+            if (p == NULL) {
+                // TODO
+                nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_OSError/*MP_QSTR_RuntimeError*/, "Out of memory"));
+            }
+        }
+    }
+    vstr_set_size(vstr, total_size + 1); // TODO: for \0
+    buf[total_size] = 0;
+    return mp_obj_new_str(qstr_from_str_take(buf, total_size + 1));
+}
+
 MP_DEFINE_CONST_FUN_OBJ_2(mp_stream_read_obj, stream_read);
+MP_DEFINE_CONST_FUN_OBJ_1(mp_stream_readall_obj, stream_readall);
 MP_DEFINE_CONST_FUN_OBJ_2(mp_stream_write_obj, stream_write);
diff --git a/py/stream.h b/py/stream.h
index 8a579b7621..b1c3c72786 100644
--- a/py/stream.h
+++ b/py/stream.h
@@ -1,2 +1,3 @@
 extern const mp_obj_fun_native_t mp_stream_read_obj;
+extern const mp_obj_fun_native_t mp_stream_readall_obj;
 extern const mp_obj_fun_native_t mp_stream_write_obj;
diff --git a/py/vstr.c b/py/vstr.c
index 80841b24ca..18e6d2d0a8 100644
--- a/py/vstr.c
+++ b/py/vstr.c
@@ -6,8 +6,8 @@
 // returned value is always at least 1 greater than argument
 #define ROUND_ALLOC(a) (((a) & ((~0) - 7)) + 8)
 
-void vstr_init(vstr_t *vstr) {
-    vstr->alloc = 32;
+void vstr_init(vstr_t *vstr, int alloc) {
+    vstr->alloc = alloc;
     vstr->len = 0;
     vstr->buf = m_new(char, vstr->alloc);
     if (vstr->buf == NULL) {
@@ -28,7 +28,16 @@ vstr_t *vstr_new(void) {
     if (vstr == NULL) {
         return NULL;
     }
-    vstr_init(vstr);
+    vstr_init(vstr, 32);
+    return vstr;
+}
+
+vstr_t *vstr_new_size(int alloc) {
+    vstr_t *vstr = m_new(vstr_t, 1);
+    if (vstr == NULL) {
+        return NULL;
+    }
+    vstr_init(vstr, alloc);
     return vstr;
 }
 
@@ -63,6 +72,36 @@ int vstr_len(vstr_t *vstr) {
     return vstr->len;
 }
 
+// Extend vstr strictly to by requested size, return pointer to newly added chunk
+char  *vstr_extend(vstr_t *vstr, int size) {
+    char *new_buf = m_renew(char, vstr->buf, vstr->alloc, vstr->alloc + size);
+    if (new_buf == NULL) {
+        vstr->had_error = true;
+        return NULL;
+    }
+    char *p = new_buf + vstr->alloc;
+    vstr->alloc += size;
+    vstr->buf = new_buf;
+    return p;
+}
+
+// Shrink vstr to be given size
+bool vstr_set_size(vstr_t *vstr, int size) {
+    char *new_buf = m_renew(char, vstr->buf, vstr->alloc, size);
+    if (new_buf == NULL) {
+        vstr->had_error = true;
+        return false;
+    }
+    vstr->buf = new_buf;
+    vstr->alloc = vstr->len = size;
+    return true;
+}
+
+// Shrink vstr allocation to its actual length
+bool vstr_shrink(vstr_t *vstr) {
+    return vstr_set_size(vstr, vstr->len);
+}
+
 bool vstr_ensure_extra(vstr_t *vstr, int size) {
     if (vstr->len + size + 1 > vstr->alloc) {
         int new_alloc = ROUND_ALLOC((vstr->len + size + 1) * 2);