py: str.join can now take arbitrary iterable as argument.

dpgeorge · dpgeorge · commit a157e4caba25 · 2014-04-09T19:17:53.000+01:00
diff --git a/py/objstr.c b/py/objstr.c
@@ -329,17 +329,19 @@ STATIC mp_obj_t str_join(mp_obj_t self_in, mp_obj_t arg) {
     mp_obj_t *seq_items;
     if (MP_OBJ_IS_TYPE(arg, &mp_type_tuple)) {
         mp_obj_tuple_get(arg, &seq_len, &seq_items);
-    } else if (MP_OBJ_IS_TYPE(arg, &mp_type_list)) {
-        mp_obj_list_get(arg, &seq_len, &seq_items);
     } else {
-        goto bad_arg;
+        if (!MP_OBJ_IS_TYPE(arg, &mp_type_list)) {
+            // arg is not a list, try to convert it to one
+            arg = mp_type_list.make_new((mp_obj_t)&mp_type_list, 1, 0, &arg);
+        }
+        mp_obj_list_get(arg, &seq_len, &seq_items);
     }
 
     // count required length
     int required_len = 0;
     for (int i = 0; i < seq_len; i++) {
         if (!MP_OBJ_IS_STR(seq_items[i])) {
-            goto bad_arg;
+            nlr_raise(mp_obj_new_exception_msg(&mp_type_TypeError, "join expected a list of str's"));
         }
         if (i > 0) {
             required_len += sep_len;
@@ -363,9 +365,6 @@ STATIC mp_obj_t str_join(mp_obj_t self_in, mp_obj_t arg) {
 
     // return joined string
     return mp_obj_str_builder_end(joined_str);
-
-bad_arg:
-    nlr_raise(mp_obj_new_exception_msg(&mp_type_TypeError, "?str.join expecting a list of str's"));
 }
 
 #define is_ws(c) ((c) == ' ' || (c) == '\t')
diff --git a/tests/basics/string-join.py b/tests/basics/string-join.py
@@ -0,0 +1,12 @@
+print(','.join(()))
+print(','.join(('a',)))
+print(','.join(('a', 'b')))
+
+print(','.join([]))
+print(','.join(['a']))
+print(','.join(['a', 'b']))
+
+print(''.join(''))
+print(''.join('abc'))
+print(','.join('abc'))
+print(','.join('abc' for i in range(5)))
