[pypy-commit] pypy default: Issue #1849: fix for str.split(), str.rsplit()

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r72918:b9445a658af8
Date: 2014-08-20 09:31 +0200
http://bitbucket.org/pypy/pypy/changeset/b9445a658af8/

Log:	Issue #1849: fix for str.split(), str.rsplit()

diff --git a/pypy/objspace/std/bytesobject.py b/pypy/objspace/std/bytesobject.py
--- a/pypy/objspace/std/bytesobject.py
+++ b/pypy/objspace/std/bytesobject.py
@@ -731,6 +731,22 @@
             return space.wrap(self._val(space).join(l))
         return self._StringMethods_descr_join(space, w_list)
 
+    _StringMethods_descr_split = descr_split
+    @unwrap_spec(maxsplit=int)
+    def descr_split(self, space, w_sep=None, maxsplit=-1):
+        if w_sep is not None and space.isinstance_w(w_sep, space.w_unicode):
+            self_as_uni = unicode_from_encoded_object(space, self, None, None)
+            return self_as_uni.descr_split(space, w_sep, maxsplit)
+        return self._StringMethods_descr_split(space, w_sep, maxsplit)
+
+    _StringMethods_descr_rsplit = descr_rsplit
+    @unwrap_spec(maxsplit=int)
+    def descr_rsplit(self, space, w_sep=None, maxsplit=-1):
+        if w_sep is not None and space.isinstance_w(w_sep, space.w_unicode):
+            self_as_uni = unicode_from_encoded_object(space, self, None, None)
+            return self_as_uni.descr_rsplit(space, w_sep, maxsplit)
+        return self._StringMethods_descr_rsplit(space, w_sep, maxsplit)
+
     def _join_return_one(self, space, w_obj):
         return (space.is_w(space.type(w_obj), space.w_str) or
                 space.is_w(space.type(w_obj), space.w_unicode))
diff --git a/pypy/objspace/std/test/test_unicodeobject.py b/pypy/objspace/std/test/test_unicodeobject.py
--- a/pypy/objspace/std/test/test_unicodeobject.py
+++ b/pypy/objspace/std/test/test_unicodeobject.py
@@ -175,6 +175,24 @@
         assert u'  a b c  '.rsplit(None, 0) == [u'  a b c']
         assert u''.rsplit('aaa') == [u'']
 
+    def test_split_rsplit_str_unicode(self):
+        x = 'abc'.split(u'b')
+        assert x == [u'a', u'c']
+        assert map(type, x) == [unicode, unicode]
+        x = 'abc'.rsplit(u'b')
+        assert x == [u'a', u'c']
+        assert map(type, x) == [unicode, unicode]
+        x = 'abc'.split(u'\u4321')
+        assert x == [u'abc']
+        assert map(type, x) == [unicode]
+        x = 'abc'.rsplit(u'\u4321')
+        assert x == [u'abc']
+        assert map(type, x) == [unicode]
+        raises(UnicodeDecodeError, '\x80'.split, u'a')
+        raises(UnicodeDecodeError, '\x80'.split, u'')
+        raises(UnicodeDecodeError, '\x80'.rsplit, u'a')
+        raises(UnicodeDecodeError, '\x80'.rsplit, u'')
+
     def test_center(self):
         s=u"a b"
         assert s.center(0) == u"a b"