From a47b64ae2dd81e45da082c6337aaf5fcef4da75c Mon Sep 17 00:00:00 2001
From: Paul Sokolovsky <pfalcon@users.sourceforge.net>
Date: Thu, 15 May 2014 07:28:19 +0300
Subject: [PATCH] objstringio: Implement io.BytesIO.

Done in generalized manner, allowing any stream class to be specified as
working with bytes.
---
 py/modio.c       |  4 +++-
 py/mpconfig.h    |  4 ++++
 py/obj.h         |  2 ++
 py/objstr.c      |  4 ++--
 py/objstr.h      |  1 +
 py/objstringio.c | 33 ++++++++++++++++++++++++++++-----
 py/stream.c      |  9 ++++++---
 7 files changed, 46 insertions(+), 11 deletions(-)

diff --git a/py/modio.c b/py/modio.c
index 892d493c9..3a009a7d8 100644
--- a/py/modio.c
+++ b/py/modio.c
@@ -37,8 +37,10 @@ STATIC const mp_map_elem_t mp_module_io_globals_table[] = {
     // Note: mp_builtin_open_obj should be defined by port, it's not
     // part of the core.
     { MP_OBJ_NEW_QSTR(MP_QSTR_open), (mp_obj_t)&mp_builtin_open_obj },
-    { MP_OBJ_NEW_QSTR(MP_QSTR_BytesIO), (mp_obj_t)&mp_type_stringio },
     { MP_OBJ_NEW_QSTR(MP_QSTR_StringIO), (mp_obj_t)&mp_type_stringio },
+    #if MICROPY_IO_BYTESIO
+    { MP_OBJ_NEW_QSTR(MP_QSTR_BytesIO), (mp_obj_t)&mp_type_bytesio },
+    #endif
 };
 
 STATIC const mp_obj_dict_t mp_module_io_globals = {
diff --git a/py/mpconfig.h b/py/mpconfig.h
index 76cd1091c..53429f1c1 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -225,6 +225,10 @@ typedef double mp_float_t;
 #define MICROPY_ENABLE_MOD_IO (1)
 #endif
 
+#ifndef MICROPY_IO_BYTESIO
+#define MICROPY_IO_BYTESIO (1)
+#endif
+
 // Whether to provide "struct" module
 #ifndef MICROPY_ENABLE_MOD_STRUCT
 #define MICROPY_ENABLE_MOD_STRUCT (1)
diff --git a/py/obj.h b/py/obj.h
index b34581c8b..e29f3efea 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -246,6 +246,7 @@ typedef struct _mp_stream_p_t {
     machine_int_t (*read)(mp_obj_t obj, void *buf, machine_uint_t size, int *errcode);
     machine_int_t (*write)(mp_obj_t obj, const void *buf, machine_uint_t size, int *errcode);
     // add seek() ?
+    int is_bytes : 1;
 } mp_stream_p_t;
 
 struct _mp_obj_type_t {
@@ -321,6 +322,7 @@ extern const mp_obj_type_t mp_type_staticmethod;
 extern const mp_obj_type_t mp_type_classmethod;
 extern const mp_obj_type_t mp_type_property;
 extern const mp_obj_type_t mp_type_stringio;
+extern const mp_obj_type_t mp_type_bytesio;
 
 // Exceptions
 extern const mp_obj_type_t mp_type_BaseException;
diff --git a/py/objstr.c b/py/objstr.c
index b5f40551f..1448f0b95 100644
--- a/py/objstr.c
+++ b/py/objstr.c
@@ -54,7 +54,7 @@ const mp_obj_t mp_const_empty_bytes;
 
 STATIC mp_obj_t mp_obj_new_str_iterator(mp_obj_t str);
 STATIC mp_obj_t mp_obj_new_bytes_iterator(mp_obj_t str);
-STATIC mp_obj_t str_new(const mp_obj_type_t *type, const byte* data, uint len);
+mp_obj_t str_new(const mp_obj_type_t *type, const byte* data, uint len);
 STATIC NORETURN void bad_implicit_conversion(mp_obj_t self_in);
 STATIC NORETURN void arg_type_mixup();
 
@@ -1612,7 +1612,7 @@ mp_obj_t mp_obj_str_builder_end(mp_obj_t o_in) {
     return o;
 }
 
-STATIC mp_obj_t str_new(const mp_obj_type_t *type, const byte* data, uint len) {
+mp_obj_t str_new(const mp_obj_type_t *type, const byte* data, uint len) {
     mp_obj_str_t *o = m_new_obj(mp_obj_str_t);
     o->base.type = type;
     o->len = len;
diff --git a/py/objstr.h b/py/objstr.h
index 019f56f7b..c0ef701a8 100644
--- a/py/objstr.h
+++ b/py/objstr.h
@@ -36,3 +36,4 @@ typedef struct _mp_obj_str_t {
 #define MP_DEFINE_STR_OBJ(obj_name, str) mp_obj_str_t obj_name = {{&mp_type_str}, 0, sizeof(str) - 1, (const byte*)str};
 
 mp_obj_t mp_obj_str_format(uint n_args, const mp_obj_t *args);
+mp_obj_t str_new(const mp_obj_type_t *type, const byte* data, uint len);
diff --git a/py/objstringio.c b/py/objstringio.c
index 880d4e117..a1a765f89 100644
--- a/py/objstringio.c
+++ b/py/objstringio.c
@@ -4,6 +4,7 @@
  * The MIT License (MIT)
  *
  * Copyright (c) 2013, 2014 Damien P. George
+ * Copyright (c) 2014 Paul Sokolovsky
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -34,6 +35,7 @@
 #include "obj.h"
 #include "runtime.h"
 #include "stream.h"
+#include "objstr.h"
 
 #if MICROPY_ENABLE_MOD_IO
 
@@ -46,7 +48,7 @@ typedef struct _mp_obj_stringio_t {
 
 STATIC void stringio_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in, mp_print_kind_t kind) {
     mp_obj_stringio_t *self = self_in;
-    print(env, "<io.StringIO 0x%x>", self->vstr);
+    print(env, self->base.type == &mp_type_stringio ? "<io.StringIO 0x%x>" : "<io.BytesIO 0x%x>", self->vstr);
 }
 
 STATIC machine_int_t stringio_read(mp_obj_t o_in, void *buf, machine_uint_t size, int *errcode) {
@@ -77,9 +79,11 @@ STATIC machine_int_t stringio_write(mp_obj_t o_in, const void *buf, machine_uint
     return size;
 }
 
+#define STREAM_TO_CONTENT_TYPE(o) (((o)->base.type == &mp_type_stringio) ? &mp_type_str : &mp_type_bytes)
+
 STATIC mp_obj_t stringio_getvalue(mp_obj_t self_in) {
     mp_obj_stringio_t *self = self_in;
-    return mp_obj_new_str((byte*)self->vstr->buf, self->vstr->len, false);
+    return str_new(STREAM_TO_CONTENT_TYPE(self), (byte*)self->vstr->buf, self->vstr->len);
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_1(stringio_getvalue_obj, stringio_getvalue);
 
@@ -96,16 +100,16 @@ mp_obj_t stringio___exit__(uint 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() {
+STATIC mp_obj_stringio_t *stringio_new(mp_obj_t type_in) {
     mp_obj_stringio_t *o = m_new_obj(mp_obj_stringio_t);
-    o->base.type = &mp_type_stringio;
+    o->base.type = type_in;
     o->vstr = vstr_new();
     o->pos = 0;
     return o;
 }
 
 STATIC mp_obj_t stringio_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_obj_t *args) {
-    mp_obj_stringio_t *o = stringio_new();
+    mp_obj_stringio_t *o = stringio_new(type_in);
 
     if (n_args > 0) {
         mp_buffer_info_t bufinfo;
@@ -135,6 +139,12 @@ STATIC const mp_stream_p_t stringio_stream_p = {
     .write = stringio_write,
 };
 
+STATIC const mp_stream_p_t bytesio_stream_p = {
+    .read = stringio_read,
+    .write = stringio_write,
+    .is_bytes = true,
+};
+
 const mp_obj_type_t mp_type_stringio = {
     { &mp_type_type },
     .name = MP_QSTR_StringIO,
@@ -146,4 +156,17 @@ const mp_obj_type_t mp_type_stringio = {
     .locals_dict = (mp_obj_t)&stringio_locals_dict,
 };
 
+#if MICROPY_IO_BYTESIO
+const mp_obj_type_t mp_type_bytesio = {
+    { &mp_type_type },
+    .name = MP_QSTR_BytesIO,
+    .print = stringio_print,
+    .make_new = stringio_make_new,
+    .getiter = mp_identity,
+    .iternext = mp_stream_unbuffered_iter,
+    .stream_p = &bytesio_stream_p,
+    .locals_dict = (mp_obj_t)&stringio_locals_dict,
+};
+#endif
+
 #endif
diff --git a/py/stream.c b/py/stream.c
index 599582b51..b1a64cdcf 100644
--- a/py/stream.c
+++ b/py/stream.c
@@ -32,6 +32,7 @@
 #include "misc.h"
 #include "qstr.h"
 #include "obj.h"
+#include "objstr.h"
 #include "stream.h"
 #if MICROPY_STREAMS_NON_BLOCK
 #include <errno.h>
@@ -53,6 +54,8 @@ STATIC mp_obj_t stream_readall(mp_obj_t self_in);
 #define is_nonblocking_error(errno) (0)
 #endif
 
+#define STREAM_CONTENT_TYPE(stream) (((stream)->is_bytes) ? &mp_type_bytes : &mp_type_str)
+
 STATIC mp_obj_t stream_read(uint n_args, const mp_obj_t *args) {
     struct _mp_obj_base_t *o = (struct _mp_obj_base_t *)args[0];
     if (o->type->stream_p == NULL || o->type->stream_p->read == NULL) {
@@ -78,7 +81,7 @@ STATIC mp_obj_t stream_read(uint n_args, const mp_obj_t *args) {
         }
         nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_OSError, "[Errno %d]", error));
     } else {
-        mp_obj_t s = mp_obj_new_str(buf, out_sz, false); // will reallocate to use exact size
+        mp_obj_t s = str_new(STREAM_CONTENT_TYPE(o->type->stream_p), buf, out_sz); // will reallocate to use exact size
         m_free(buf, sz);
         return s;
     }
@@ -155,7 +158,7 @@ STATIC mp_obj_t stream_readall(mp_obj_t self_in) {
         }
     }
 
-    mp_obj_t s = mp_obj_new_str((byte*)vstr->buf, total_size, false);
+    mp_obj_t s = str_new(STREAM_CONTENT_TYPE(o->type->stream_p), (byte*)vstr->buf, total_size);
     vstr_free(vstr);
     return s;
 }
@@ -204,7 +207,7 @@ STATIC mp_obj_t stream_unbuffered_readline(uint n_args, const mp_obj_t *args) {
         }
     }
     // TODO need a string creation API that doesn't copy the given data
-    mp_obj_t ret = mp_obj_new_str((byte*)vstr->buf, vstr->len, false);
+    mp_obj_t ret = str_new(STREAM_CONTENT_TYPE(o->type->stream_p), (byte*)vstr->buf, vstr->len);
     vstr_free(vstr);
     return ret;
 }
-- 
GitLab