[pypy-commit] pypy py3k: improve bytearray's starts/endswith error messages too

[pjenvey]

Author: Philip Jenvey <pjenvey at underboss.org>
Branch: py3k
Changeset: r60114:d50054f9d114
Date: 2013-01-16 16:09 -0800
http://bitbucket.org/pypy/pypy/changeset/d50054f9d114/

Log:	improve bytearray's starts/endswith error messages too

diff --git a/pypy/objspace/std/bytearrayobject.py b/pypy/objspace/std/bytearrayobject.py
--- a/pypy/objspace/std/bytearrayobject.py
+++ b/pypy/objspace/std/bytearrayobject.py
@@ -261,8 +261,18 @@
     return stringobject.str_rfind__String_String_ANY_ANY(space, w_str, w_char,
                                                          w_start, w_stop)
 
+def _suffix_to_str(space, w_suffix, funcname):
+    try:
+        return space.bufferstr_new_w(w_suffix)
+    except OperationError as e:
+        if e.match(space, space.w_TypeError):
+            msg = ("%s first arg must be bytes or a tuple of bytes, "
+                   "not %s")
+            typename = space.type(w_suffix).getname(space)
+            raise operationerrfmt(space.w_TypeError, msg, funcname, typename)
+
 def str_startswith__Bytearray_ANY_ANY_ANY(space, w_bytearray, w_prefix, w_start, w_stop):
-    w_prefix = space.wrapbytes(space.bufferstr_new_w(w_prefix))
+    w_prefix = space.wrapbytes(_suffix_to_str(space, w_prefix, 'startswith'))
     w_str = _to_bytes(space, w_bytearray)
     return stringobject.str_startswith__String_String_ANY_ANY(space, w_str, w_prefix,
                                                               w_start, w_stop)
@@ -275,7 +285,7 @@
                                                               w_start, w_stop)
 
 def str_endswith__Bytearray_ANY_ANY_ANY(space, w_bytearray, w_suffix, w_start, w_stop):
-    w_suffix = space.wrapbytes(space.bufferstr_new_w(w_suffix))
+    w_suffix = space.wrapbytes(_suffix_to_str(space, w_suffix, 'endswith'))
     w_str = _to_bytes(space, w_bytearray)
     return stringobject.str_endswith__String_String_ANY_ANY(space, w_str, w_suffix,
                                                               w_start, w_stop)
diff --git a/pypy/objspace/std/test/test_bytearrayobject.py b/pypy/objspace/std/test/test_bytearrayobject.py
--- a/pypy/objspace/std/test/test_bytearrayobject.py
+++ b/pypy/objspace/std/test/test_bytearrayobject.py
@@ -183,6 +183,18 @@
         assert bytearray(b'hello').endswith(b'lo')
         assert bytearray(b'hello').endswith(bytearray(b'lo'))
         assert bytearray(b'hello').endswith((bytearray(b'lo'), b'he'))
+        try:
+            bytearray(b'hello').startswith([b'o'])
+        except TypeError as e:
+            assert 'bytes' in str(e)
+        else:
+            assert False, 'Expected TypeError'
+        try:
+            bytearray(b'hello').endswith([b'o'])
+        except TypeError as e:
+            assert 'bytes' in str(e)
+        else:
+            assert False, 'Expected TypeError'
 
     def test_stringlike_conversions(self):
         # methods that should return bytearray (and not str)