[pypy-svn] r77896 - in pypy/branch/fast-forward/pypy/objspace/std: . test

[afa at codespeak.net]

Author: afa
Date: Thu Oct 14 08:52:06 2010
New Revision: 77896

Modified:
   pypy/branch/fast-forward/pypy/objspace/std/bytearraytype.py
   pypy/branch/fast-forward/pypy/objspace/std/test/test_bytes.py
Log:
Enhance bytearray constructor


Modified: pypy/branch/fast-forward/pypy/objspace/std/bytearraytype.py
==============================================================================
--- pypy/branch/fast-forward/pypy/objspace/std/bytearraytype.py	(original)
+++ pypy/branch/fast-forward/pypy/objspace/std/bytearraytype.py	Thu Oct 14 08:52:06 2010
@@ -1,6 +1,7 @@
 import sys
 from pypy.interpreter import gateway
 from pypy.interpreter.baseobjspace import ObjSpace, W_Root
+from pypy.interpreter.error import OperationError
 from pypy.objspace.std.register_all import register_all
 from pypy.objspace.std.stdtypedef import StdTypeDef, SMM
 
@@ -14,15 +15,66 @@
     str_join, str_split, str_rsplit, str_partition, str_rpartition,
     str_splitlines)
 
- at gateway.unwrap_spec(ObjSpace, W_Root, W_Root, W_Root, W_Root)
-def descr__new__(space, w_bytearraytype,
-                 w_source='', w_encoding=None, w_errors=None):
+def _getbytevalue(space, w_value):
+    if space.isinstance_w(w_value, space.w_str):
+        string = space.str_w(w_value)
+        if len(string) != 1:
+            raise OperationError(space.w_ValueError, space.wrap(
+                "string must be of size 1"))
+        return string[0]
+
+    value = space.getindex_w(w_value, None)
+    if not 0 <= value < 256:
+        # this includes the OverflowError in case the long is too large
+        raise OperationError(space.w_ValueError, space.wrap(
+            "byte must be in range(0, 256)"))
+    return chr(value)
+
+def new_bytearray(space, w_bytearraytype, data):
     from pypy.objspace.std.bytearrayobject import W_BytearrayObject
-    data = [c for c in space.str_w(w_source)]
     w_obj = space.allocate_instance(W_BytearrayObject, w_bytearraytype)
     W_BytearrayObject.__init__(w_obj, data)
     return w_obj
 
+ at gateway.unwrap_spec(ObjSpace, W_Root, W_Root, W_Root, W_Root)
+def descr__new__(space, w_bytearraytype,
+                 w_source='', w_encoding=None, w_errors=None):
+
+    data = []
+    # String-like argument
+    try:
+        string = space.str_w(w_source)
+    except OperationError, e:
+        if not e.match(space, space.w_TypeError):
+            raise
+    else:
+        data = [c for c in string]
+        return new_bytearray(space, w_bytearraytype, data)
+
+    # Is it an int?
+    try:
+        count = space.int_w(w_source)
+    except OperationError, e:
+        if not e.match(space, space.w_TypeError):
+            raise
+    else:
+        data = ['\0'] * count
+        return new_bytearray(space, w_bytearraytype, data)
+
+    # sequence of bytes
+    w_iter = space.iter(w_source)
+    while True:
+        try:
+            w_item = space.next(w_iter)
+        except OperationError, e:
+            if not e.match(space, space.w_StopIteration):
+                raise
+            break
+        value = _getbytevalue(space, w_item)
+        data.append(value)
+
+    return new_bytearray(space, w_bytearraytype, data)
+
 # ____________________________________________________________
 
 bytearray_typedef = StdTypeDef("bytearray",

Modified: pypy/branch/fast-forward/pypy/objspace/std/test/test_bytes.py
==============================================================================
--- pypy/branch/fast-forward/pypy/objspace/std/test/test_bytes.py	(original)
+++ pypy/branch/fast-forward/pypy/objspace/std/test/test_bytes.py	Thu Oct 14 08:52:06 2010
@@ -5,6 +5,16 @@
         assert type(b) is bytearray
         assert b.__class__ is bytearray
 
+    def test_constructor(self):
+        assert bytearray() == ""
+        assert bytearray('abc') == "abc"
+        assert bytearray(['a', 'b', 'c']) == "abc"
+        assert bytearray([65, 66, 67]) == "ABC"
+        assert bytearray(5) == '\0' * 5
+        raises(ValueError, bytearray, ['a', 'bc'])
+        raises(ValueError, bytearray, [65, -3])
+        raises(TypeError, bytearray, [65.0])
+
     def test_len(self):
         b = bytearray('test')
         assert len(b) == 4