Merge remote-tracking branch 'upstream/master' into builtins

14 files changed, 226 insertions, 93 deletions

diff --git a/py/builtin.c b/py/builtin.c
index 39b2da9200..8f93e843b5 100644
--- a/py/builtin.c
+++ b/py/builtin.c
@@ -339,7 +339,7 @@ static mp_obj_t mp_builtin_sorted(mp_obj_t args, mp_map_t *kwargs) {
     }
     mp_obj_t self = list_type.make_new((mp_obj_t)&list_type, 1, args_items);
     mp_obj_t new_args = rt_build_tuple(1, &self);
-    list_sort(new_args, kwargs);
+    mp_obj_list_sort(new_args, kwargs);
 
     return self;
 }
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/obj.h b/py/obj.h
index 6a6f98d0cd..00d152ec4b 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -292,7 +292,7 @@ extern const mp_obj_type_t list_type;
 mp_obj_t mp_obj_list_append(mp_obj_t self_in, mp_obj_t arg);
 void mp_obj_list_get(mp_obj_t self_in, uint *len, mp_obj_t **items);
 void mp_obj_list_store(mp_obj_t self_in, mp_obj_t index, mp_obj_t value);
-mp_obj_t list_sort(mp_obj_t args, struct _mp_map_t *kwargs);
+mp_obj_t mp_obj_list_sort(mp_obj_t args, struct _mp_map_t *kwargs);
 
 // enumerate
 extern const mp_obj_type_t enumerate_type;
diff --git a/py/objdict.c b/py/objdict.c
index 6dbb1f316b..e164ed74ca 100644
--- a/py/objdict.c
+++ b/py/objdict.c
@@ -57,6 +57,12 @@ static mp_obj_t dict_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
                 return elem->value;
             }
         }
+        case RT_COMPARE_OP_IN:
+        case RT_COMPARE_OP_NOT_IN:
+        {
+            mp_map_elem_t *elem = mp_map_lookup(&o->map, rhs_in, MP_MAP_LOOKUP);
+            return MP_BOOL((op == RT_COMPARE_OP_IN) ^ (elem == NULL));
+        }
         default:
             // op not supported
             return NULL;
@@ -362,10 +368,20 @@ static void dict_view_print(void (*print)(void *env, const char *fmt, ...), void
     print(env, "])");
 }
 
+static mp_obj_t dict_view_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
+    /* only supported for the 'keys' kind until sets and dicts are refactored */
+    mp_obj_dict_view_t *o = lhs_in;
+    if (o->kind != MP_DICT_VIEW_KEYS) return NULL;
+    if (op != RT_COMPARE_OP_IN && op != RT_COMPARE_OP_NOT_IN) return NULL;
+    return dict_binary_op(op, o->dict, rhs_in);
+}
+
+
 static const mp_obj_type_t dict_view_type = {
     { &mp_const_type },
     "dict_view",
     .print = dict_view_print,
+    .binary_op = dict_view_binary_op,
     .getiter = dict_view_getiter,
 };
 
diff --git a/py/objexcept.c b/py/objexcept.c
index 4708a27bfc..f083e61e52 100644
--- a/py/objexcept.c
+++ b/py/objexcept.c
@@ -8,105 +8,86 @@
 #include "misc.h"
 #include "mpconfig.h"
 #include "obj.h"
+#include "objtuple.h"
 
+// This is unified class for C-level and Python-level exceptions
+// Python-level exception have empty ->msg and all arguments are in
+// args tuple. C-level excepttion likely have ->msg, and may as well
+// have args tuple (or otherwise have it as NULL).
 typedef struct mp_obj_exception_t {
     mp_obj_base_t base;
     qstr id;
-    int n_args;
-    const void *args[];
+    qstr msg;
+    mp_obj_tuple_t args;
 } mp_obj_exception_t;
 
 void exception_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t o_in) {
     mp_obj_exception_t *o = o_in;
-    switch (o->n_args) {
-        case 0:
-            print(env, "%s", qstr_str(o->id));
-            break;
-        case 1:
-            print(env, "%s: %s", qstr_str(o->id), (const char*)o->args[0]);
-            break;
-        case 2:
-            print(env, "%s: ", qstr_str(o->id));
-            print(env, (const char*)o->args[0], o->args[1]);
-            break;
-        default: // here we just assume at least 3 args, but only use first 3
-            print(env, "%s: ", qstr_str(o->id));
-            print(env, (const char*)o->args[0], o->args[1], o->args[2]);
-            break;
+    if (o->msg != 0) {
+        print(env, "%s: %s", qstr_str(o->id), qstr_str(o->msg));
+    } else {
+        print(env, "%s", qstr_str(o->id));
+        tuple_print(print, env, &o->args);
     }
 }
 
+// args in reversed order
+static mp_obj_t exception_call(mp_obj_t self_in, int n_args, const mp_obj_t *args) {
+    mp_obj_exception_t *base = self_in;
+    mp_obj_exception_t *o = m_new_obj_var(mp_obj_exception_t, mp_obj_t*, n_args);
+    o->base.type = &exception_type;
+    o->id = base->id;
+    o->msg = 0;
+    o->args.len = n_args;
+
+    // TODO: factor out as reusable copy_reversed()
+    int j = 0;
+    for (int i = n_args - 1; i >= 0; i--) {
+        o->args.items[i] = args[j++];
+    }
+    return o;
+}
+
 const mp_obj_type_t exception_type = {
     { &mp_const_type },
     "exception",
     .print = exception_print,
+    .call_n = exception_call,
 };
 
 mp_obj_t mp_obj_new_exception(qstr id) {
-    mp_obj_exception_t *o = m_new_obj(mp_obj_exception_t);
-    o->base.type = &exception_type;
-    o->id = id;
-    o->n_args = 0;
-    return o;
+    return mp_obj_new_exception_msg_varg(id, NULL);
 }
 
 mp_obj_t mp_obj_new_exception_msg(qstr id, const char *msg) {
-    mp_obj_exception_t *o = m_new_obj_var(mp_obj_exception_t, void*, 1);
-    o->base.type = &exception_type;
-    o->id = id;
-    o->n_args = 1;
-    o->args[0] = msg;
-    return o;
+    return mp_obj_new_exception_msg_varg(id, msg);
 }
 
 mp_obj_t mp_obj_new_exception_msg_1_arg(qstr id, const char *fmt, const char *a1) {
-    mp_obj_exception_t *o = m_new_obj_var(mp_obj_exception_t, void*, 2);
-    o->base.type = &exception_type;
-    o->id = id;
-    o->n_args = 2;
-    o->args[0] = fmt;
-    o->args[1] = a1;
-    return o;
+    return mp_obj_new_exception_msg_varg(id, fmt, a1);
 }
 
 mp_obj_t mp_obj_new_exception_msg_2_args(qstr id, const char *fmt, const char *a1, const char *a2) {
-    mp_obj_exception_t *o = m_new_obj_var(mp_obj_exception_t, void*, 3);
-    o->base.type = &exception_type;
-    o->id = id;
-    o->n_args = 3;
-    o->args[0] = fmt;
-    o->args[1] = a1;
-    o->args[2] = a2;
-    return o;
+    return mp_obj_new_exception_msg_varg(id, fmt, a1, a2);
 }
 
 mp_obj_t mp_obj_new_exception_msg_varg(qstr id, const char *fmt, ...) {
-    // count number of arguments by number of % signs, excluding %%
-    int n_args = 1; // count fmt
-    for (const char *s = fmt; *s; s++) {
-        if (*s == '%') {
-            if (s[1] == '%') {
-                s += 1;
-            } else {
-                n_args += 1;
-            }
-        }
-    }
-
     // make exception object
-    mp_obj_exception_t *o = m_new_obj_var(mp_obj_exception_t, void*, n_args);
+    mp_obj_exception_t *o = m_new_obj_var(mp_obj_exception_t, mp_obj_t*, 0);
     o->base.type = &exception_type;
     o->id = id;
-    o->n_args = n_args;
-    o->args[0] = fmt;
-
-    // extract args and store them
-    va_list ap;
-    va_start(ap, fmt);
-    for (int i = 1; i < n_args; i++) {
-        o->args[i] = va_arg(ap, void*);
+    o->args.len = 0;
+    if (fmt == NULL) {
+        o->msg = 0;
+    } else {
+        // render exception message
+        vstr_t *vstr = vstr_new();
+        va_list ap;
+        va_start(ap, fmt);
+        vstr_vprintf(vstr, fmt, ap);
+        va_end(ap);
+        o->msg = qstr_from_str_take(vstr->buf, vstr->alloc);
     }
-    va_end(ap);
 
     return o;
 }
diff --git a/py/objlist.c b/py/objlist.c
index f806dfae8f..100a02f3cf 100644
--- a/py/objlist.c
+++ b/py/objlist.c
@@ -248,7 +248,7 @@ static void mp_quicksort(mp_obj_t *head, mp_obj_t *tail, mp_obj_t key_fn, bool r
     }
 }
 
-mp_obj_t list_sort(mp_obj_t args, mp_map_t *kwargs) {
+mp_obj_t mp_obj_list_sort(mp_obj_t args, mp_map_t *kwargs) {
     mp_obj_t *args_items = NULL;
     uint args_len = 0;
 
@@ -381,7 +381,7 @@ static MP_DEFINE_CONST_FUN_OBJ_3(list_insert_obj, list_insert);
 static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(list_pop_obj, 1, 2, list_pop);
 static MP_DEFINE_CONST_FUN_OBJ_2(list_remove_obj, list_remove);
 static MP_DEFINE_CONST_FUN_OBJ_1(list_reverse_obj, list_reverse);
-static MP_DEFINE_CONST_FUN_OBJ_KW(list_sort_obj, 0, list_sort);
+static MP_DEFINE_CONST_FUN_OBJ_KW(list_sort_obj, 0, mp_obj_list_sort);
 
 static const mp_method_t list_type_methods[] = {
     { "append", &list_append_obj },
diff --git a/py/objset.c b/py/objset.c
index e41f2c47f4..ba083ada8a 100644
--- a/py/objset.c
+++ b/py/objset.c
@@ -45,6 +45,7 @@ void set_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj
     print(env, "}");
 }
 
+
 static mp_obj_t set_make_new(mp_obj_t type_in, int n_args, const mp_obj_t *args) {
     switch (n_args) {
         case 0:
@@ -405,6 +406,13 @@ static mp_obj_t set_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
         return set_issuperset(lhs, rhs);
     case RT_COMPARE_OP_NOT_EQUAL:
         return MP_BOOL(set_equal(lhs, rhs) == mp_const_false);
+    case RT_COMPARE_OP_IN:
+    case RT_COMPARE_OP_NOT_IN:
+    {
+        mp_obj_set_t *o = lhs;
+        mp_obj_t elem = mp_set_lookup(&o->set, rhs, MP_MAP_LOOKUP);
+        return MP_BOOL((op == RT_COMPARE_OP_IN) ^ (elem == NULL));
+    }
     default:
         // op not supported
         return NULL;
diff --git a/py/objstr.c b/py/objstr.c
index f48bde6001..8b7ab9692f 100644
--- a/py/objstr.c
+++ b/py/objstr.c
@@ -85,6 +85,15 @@ mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
                 return mp_obj_new_str(qstr_from_str_take(val, alloc_len));
             }
             break;
+        case RT_COMPARE_OP_IN:
+        case RT_COMPARE_OP_NOT_IN:
+            /* NOTE `a in b` is `b.__contains__(a)` */
+            if (MP_OBJ_IS_TYPE(rhs_in, &str_type)) {
+                const char *rhs_str = qstr_str(((mp_obj_str_t*)rhs_in)->qstr);
+                /* FIXME \0 in strs */
+                return MP_BOOL((op == RT_COMPARE_OP_IN) ^ (strstr(lhs_str, rhs_str) == NULL));
+            }
+            break;
     }
 
     return MP_OBJ_NULL; // op not supported
diff --git a/py/objzip.c b/py/objzip.c
index 2ab80af0b5..a552ff5881 100644
--- a/py/objzip.c
+++ b/py/objzip.c
@@ -12,12 +12,6 @@ typedef struct _mp_obj_zip_t {
     mp_obj_t iters[];
 } mp_obj_zip_t;
 
-static mp_obj_t zip_getiter(mp_obj_t self_in) {
-    return self_in;
-}
-
-static mp_obj_t zip_iternext(mp_obj_t self_in);
-
 static mp_obj_t zip_make_new(mp_obj_t type_in, int n_args, const mp_obj_t *args) {
     /* NOTE: args are backwards */
     mp_obj_zip_t *o = m_new_obj_var(mp_obj_zip_t, mp_obj_t, n_args);
@@ -29,13 +23,9 @@ static mp_obj_t zip_make_new(mp_obj_t type_in, int n_args, const mp_obj_t *args)
     return o;
 }
 
-const mp_obj_type_t zip_type = {
-    { &mp_const_type },
-    "zip",
-    .make_new = zip_make_new,
-    .iternext = zip_iternext,
-    .getiter = zip_getiter,
-};
+static mp_obj_t zip_getiter(mp_obj_t self_in) {
+    return self_in;
+}
 
 static mp_obj_t zip_iternext(mp_obj_t self_in) {
     assert(MP_OBJ_IS_TYPE(self_in, &zip_type));
@@ -57,3 +47,11 @@ static mp_obj_t zip_iternext(mp_obj_t self_in) {
     }
     return o;
 }
+
+const mp_obj_type_t zip_type = {
+    { &mp_const_type },
+    "zip",
+    .make_new = zip_make_new,
+    .getiter = zip_getiter,
+    .iternext = zip_iternext,
+};
diff --git a/py/runtime.c b/py/runtime.c
index 6047b90b40..a76af3db7c 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -569,22 +569,57 @@ mp_obj_t rt_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
         } else if (MP_OBJ_IS_TYPE(rhs, &complex_type)) {
             return mp_obj_complex_binary_op(op, lhs_val, 0, rhs);
         }
-    } else {
-        if (MP_OBJ_IS_OBJ(lhs)) {
-            mp_obj_base_t *o = lhs;
+    }
+
+    /* deal with `in` and `not in`
+     *
+     * NOTE `a in b` is `b.__contains__(a)`, hence why the generic dispatch
+     * needs to go below
+     */
+    if (op == RT_COMPARE_OP_IN || op == RT_COMPARE_OP_NOT_IN) {
+        if (!MP_OBJ_IS_SMALL_INT(rhs)) {
+            mp_obj_base_t *o = rhs;
             if (o->type->binary_op != NULL) {
-                mp_obj_t result = o->type->binary_op(op, lhs, rhs);
-                if (result != NULL) {
-                    return result;
+                mp_obj_t res = o->type->binary_op(op, rhs, lhs);
+                if (res != NULL) {
+                    return res;
                 }
             }
+            if (o->type->getiter != NULL) {
+                /* second attempt, walk the iterator */
+                mp_obj_t next = NULL;
+                mp_obj_t iter = rt_getiter(rhs);
+                while ((next = rt_iternext(iter)) != mp_const_stop_iteration) {
+                    if (mp_obj_equal(next, lhs)) {
+                        return MP_BOOL(op == RT_COMPARE_OP_IN);
+                    }
+                }
+                return MP_BOOL(op != RT_COMPARE_OP_IN);
+            }
+        }
+
+        nlr_jump(mp_obj_new_exception_msg_varg(
+                     MP_QSTR_TypeError, "'%s' object is not iterable",
+                     mp_obj_get_type_str(rhs)));
+        return mp_const_none;
+    }
+
+    if (MP_OBJ_IS_OBJ(lhs)) {
+        mp_obj_base_t *o = lhs;
+        if (o->type->binary_op != NULL) {
+            mp_obj_t result = o->type->binary_op(op, lhs, rhs);
+            if (result != NULL) {
+                return result;
+            }
         }
+        // TODO implement dispatch for reverse binary ops
     }
 
     // TODO specify in error message what the operator is
     nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_TypeError,
         "unsupported operand types for binary operator: '%s', '%s'",
         mp_obj_get_type_str(lhs), mp_obj_get_type_str(rhs)));
+    return mp_const_none;
 }
 
 mp_obj_t rt_make_function_from_id(int unique_code_id) {
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);