[pypy-commit] pypy utf8-unicode2: Add support for str.join([<list of chars>])

[waedt]

Author: Tyler Wade <wayedt at gmail.com>
Branch: utf8-unicode2
Changeset: r72696:9661a2f7dd0f
Date: 2014-08-04 14:36 -0500
http://bitbucket.org/pypy/pypy/changeset/9661a2f7dd0f/

Log:	Add support for str.join([<list of chars>])

diff --git a/rpython/rtyper/lltypesystem/rstr.py b/rpython/rtyper/lltypesystem/rstr.py
--- a/rpython/rtyper/lltypesystem/rstr.py
+++ b/rpython/rtyper/lltypesystem/rstr.py
@@ -820,6 +820,32 @@
             i += 1
         return result
 
+    def ll_join_chars_with_str(s, length, chars):
+        s_chars = s.chars
+        s_len = len(s_chars)
+        num_chars = length
+        if num_chars == 0:
+            return s.empty()
+
+        try:
+            seplen = ovfcheck(s_len * (num_chars - 1))
+        except OverflowError:
+            raise MemoryError
+
+        # a single '+' at the end is allowed to overflow: it gets
+        # a negative result, and the gc will complain
+        result = s.malloc(num_chars + seplen)
+        res_index = 1
+        result.chars[0] = chars[0]
+        i = 1
+        while i < num_chars:
+            s.copy_contents(s, result, 0, res_index, s_len)
+            res_index += s_len
+            result.chars[res_index] = chars[i]
+            res_index += 1
+            i += 1
+        return result
+
     @jit.oopspec('stroruni.slice(s1, start, stop)')
     @signature(types.any(), types.int(), types.int(), returns=types.any())
     @jit.elidable
diff --git a/rpython/rtyper/rstr.py b/rpython/rtyper/rstr.py
--- a/rpython/rtyper/rstr.py
+++ b/rpython/rtyper/rstr.py
@@ -235,6 +235,8 @@
         else:
             if r_lst.item_repr == rstr.repr:
                 llfn = self.ll.ll_join
+            elif r_lst.item_repr == char_repr:
+                llfn = self.ll.ll_join_chars_with_str
             else:
                 raise TyperError("sep.join() of non-string list: %r" % r_lst)
             return hop.gendirectcall(llfn, v_str, v_length, v_items)
diff --git a/rpython/rtyper/test/test_rstr.py b/rpython/rtyper/test/test_rstr.py
--- a/rpython/rtyper/test/test_rstr.py
+++ b/rpython/rtyper/test/test_rstr.py
@@ -516,6 +516,9 @@
         res = self.interpret(lambda: const('.').join([const('abc'), const('def')]), [])
         assert self.ll_to_string(res) == const('abc.def')
 
+        res = self.interpret(lambda: const('  ').join([const('a'), const('b'), const('c')]), [])
+        assert self.ll_to_string(res) == const('a  b  c')
+
         def fn(i, j):
             s1 = [ const(''), const(','), const(' and ')]
             s2 = [ [], [const('foo')], [const('bar'), const('baz'), const('bazz')]]