Merge pull request #131 from chipaca/dict_fromkeys

Added dict.fromkeys.

py/builtin.c

@@ -166,25 +166,12 @@ static mp_obj_t mp_builtin_iter(mp_obj_t o_in) {
 MP_DEFINE_CONST_FUN_OBJ_1(mp_builtin_iter_obj, mp_builtin_iter);
 
 mp_obj_t mp_builtin_len(mp_obj_t o_in) {
-    mp_small_int_t len = 0;
-    if (MP_OBJ_IS_TYPE(o_in, &str_type)) {
-        len = strlen(qstr_str(mp_obj_str_get(o_in)));
-    } else if (MP_OBJ_IS_TYPE(o_in, &tuple_type)) {
-        uint seq_len;
-        mp_obj_t *seq_items;
-        mp_obj_tuple_get(o_in, &seq_len, &seq_items);
-        len = seq_len;
-    } else if (MP_OBJ_IS_TYPE(o_in, &list_type)) {
-        uint seq_len;
-        mp_obj_t *seq_items;
-        mp_obj_list_get(o_in, &seq_len, &seq_items);
-        len = seq_len;
-    } else if (MP_OBJ_IS_TYPE(o_in, &dict_type)) {
-        len = mp_obj_dict_len(o_in);
-    } else {
+    mp_obj_t len = mp_obj_len_maybe(o_in);
+    if (len == NULL) {
         nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_TypeError, "object of type '%s' has no len()", mp_obj_get_type_str(o_in)));
+    } else {
+        return len;
     }
-    return MP_OBJ_NEW_SMALL_INT(len);
 }
 
 mp_obj_t mp_builtin_max(int n_args, const mp_obj_t *args) {

py/obj.c

@@ -2,6 +2,7 @@
 #include <stdlib.h>
 #include <stdio.h>
 #include <stdarg.h>
+#include <string.h>
 #include <assert.h>
 
 #include "nlr.h"
@@ -229,3 +230,26 @@ uint mp_get_index(const mp_obj_type_t *type, machine_uint_t len, mp_obj_t index)
         nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_TypeError, "%s indices must be integers, not %s", type->name, mp_obj_get_type_str(index)));
     }
 }
+
+// may return NULL
+mp_obj_t mp_obj_len_maybe(mp_obj_t o_in) {
+    mp_small_int_t len = 0;
+    if (MP_OBJ_IS_TYPE(o_in, &str_type)) {
+        len = strlen(qstr_str(mp_obj_str_get(o_in)));
+    } else if (MP_OBJ_IS_TYPE(o_in, &tuple_type)) {
+        uint seq_len;
+        mp_obj_t *seq_items;
+        mp_obj_tuple_get(o_in, &seq_len, &seq_items);
+        len = seq_len;
+    } else if (MP_OBJ_IS_TYPE(o_in, &list_type)) {
+        uint seq_len;
+        mp_obj_t *seq_items;
+        mp_obj_list_get(o_in, &seq_len, &seq_items);
+        len = seq_len;
+    } else if (MP_OBJ_IS_TYPE(o_in, &dict_type)) {
+        len = mp_obj_dict_len(o_in);
+    } else {
+        return NULL;
+    }
+    return MP_OBJ_NEW_SMALL_INT(len);
+}

py/obj.h

@@ -234,6 +234,7 @@ void mp_obj_get_complex(mp_obj_t self_in, mp_float_t *real, mp_float_t *imag);
 qstr mp_obj_get_qstr(mp_obj_t arg);
 mp_obj_t *mp_obj_get_array_fixed_n(mp_obj_t o, machine_int_t n);
 uint mp_get_index(const mp_obj_type_t *type, machine_uint_t len, mp_obj_t index);
+mp_obj_t mp_obj_len_maybe(mp_obj_t o_in); /* may return NULL */
 
 // none
 extern const mp_obj_type_t none_type;

py/objdict.c

@@ -376,6 +376,61 @@ static mp_obj_t dict_values(mp_obj_t self_in) {
 }
 static MP_DEFINE_CONST_FUN_OBJ_1(dict_values_obj, dict_values);
 
+
+/******************************************************************************/
+/* dict metaclass                                                             */
+
+static mp_obj_t dict_fromkeys(int n_args, const mp_obj_t *args) {
+    assert(2 <= n_args && n_args <= 3);
+    mp_obj_t iter = rt_getiter(args[1]);
+    mp_obj_t len = mp_obj_len_maybe(iter);
+    mp_obj_t value = mp_const_none;
+    mp_obj_t next = NULL;
+    mp_obj_dict_t *self = NULL;
+
+    if (n_args > 2) {
+        value = args[2];
+    }
+
+    if (len == NULL) {
+        /* object's type doesn't have a __len__ slot */
+        self = mp_obj_new_dict(0);
+    } else {
+        self = mp_obj_new_dict(MP_OBJ_SMALL_INT_VALUE(len));
+    }
+
+    while ((next = rt_iternext(iter)) != mp_const_stop_iteration) {
+        mp_map_lookup(&self->map, next, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND)->value = value;
+    }
+
+    return self;
+}
+static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_fromkeys_obj, 2, 3, dict_fromkeys);
+
+static const mp_method_t dict_class_methods[] = {
+    { "fromkeys", &dict_fromkeys_obj },
+    { NULL, NULL }, // end-of-list sentinel
+};
+
+/* this should be unnecessary when inheritance works */
+static void dict_class_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in) {
+    print(env, "<class 'dict'>");
+}
+
+/* this should be unnecessary when inheritance works */
+static mp_obj_t dict_class_call_n(mp_obj_t self_in, int n_args, const mp_obj_t *args) {
+    return rt_build_map(0);
+}
+
+static const mp_obj_type_t dict_class = {
+    { &mp_const_type },
+    "dict_class",
+    .print = dict_class_print,
+    .methods = dict_class_methods,
+    .call_n = dict_class_call_n,
+};
+
+
 /******************************************************************************/
 /* dict constructors & etc                                                    */
 
@@ -394,7 +449,7 @@ static const mp_method_t dict_type_methods[] = {
 };
 
 const mp_obj_type_t dict_type = {
-    { &mp_const_type },
+    { &dict_class },
     "dict",
     .print = dict_print,
     .make_new = dict_make_new,

tests/basics/tests/dict_fromkeys.py

@@ -0,0 +1,10 @@
+d = dict.fromkeys([1, 2, 3, 4])
+l = list(d.keys())
+l.sort()
+print(l)
+
+d = dict.fromkeys([1, 2, 3, 4], 42)
+l = list(d.values())
+l.sort()
+print(l)
+