From a157e4caba25e61dba2cda9bac37920887e5ec02 Mon Sep 17 00:00:00 2001 From: Damien George Date: Wed, 9 Apr 2014 19:17:53 +0100 Subject: [PATCH] py: str.join can now take arbitrary iterable as argument. --- py/objstr.c | 13 ++++++------- tests/basics/string-join.py | 12 ++++++++++++ 2 files changed, 18 insertions(+), 7 deletions(-) create mode 100644 tests/basics/string-join.py diff --git a/py/objstr.c b/py/objstr.c index 5e9f0c76b6..06bc98a723 100644 --- 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 new file mode 100644 index 0000000000..275a804c64 --- /dev/null +++ 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)))