[pypy-svn] pypy fast-forward: add pickle support to BytesIO

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: fast-forward
Changeset: r40458:dc782a71d278
Date: 2011-01-07 13:04 +0100
http://bitbucket.org/pypy/pypy/changeset/dc782a71d278/

Log:	add pickle support to BytesIO

diff --git a/pypy/module/_io/test/test_bytesio.py b/pypy/module/_io/test/test_bytesio.py
--- a/pypy/module/_io/test/test_bytesio.py
+++ b/pypy/module/_io/test/test_bytesio.py
@@ -42,3 +42,16 @@
         f.seek(3)
         assert f.truncate() == 3
         assert f.getvalue() == "hel"
+
+    def test_setstate(self):
+        # state is (content, position, __dict__)
+        import _io
+        f = _io.BytesIO("hello")
+        content, pos, __dict__ = f.__getstate__()
+        assert (content, pos) == ("hello", 0)
+        assert __dict__ is None or __dict__ == {}
+        f.__setstate__(("world", 3, {"a": 1}))
+        assert f.getvalue() == "world"
+        assert f.read() == "ld"
+        assert f.a == 1
+        assert f.__getstate__() == ("world", 5, {"a": 1})

diff --git a/pypy/module/_io/interp_bytesio.py b/pypy/module/_io/interp_bytesio.py
--- a/pypy/module/_io/interp_bytesio.py
+++ b/pypy/module/_io/interp_bytesio.py
@@ -183,6 +183,24 @@
     def closed_get_w(space, self):
         return space.wrap(self.buf is None)
 
+    @unwrap_spec('self', ObjSpace)
+    def getstate_w(self, space):
+        w_content = space.wrap(buffer2string(self.buf, 0, self.string_size))
+        return space.newtuple([
+            w_content,
+            space.wrap(self.pos),
+            self.getdict()])
+
+    @unwrap_spec('self', ObjSpace, W_Root)
+    def setstate_w(self, space, w_state):
+        w_content, w_pos, w_dict = space.unpackiterable(w_state, 3)
+        pos = space.int_w(w_pos)
+        self.buf = []
+        self.write_w(space, w_content)
+        self.pos = pos
+        if not space.is_w(w_dict, space.w_None):
+            space.call_method(self.getdict(), "update", w_dict)
+
 W_BytesIO.typedef = TypeDef(
     'BytesIO', W_BufferedIOBase.typedef,
     __new__ = generic_new_descr(W_BytesIO),
@@ -201,5 +219,7 @@
     seekable = interp2app(W_BytesIO.seekable_w),
     close = interp2app(W_BytesIO.close_w),
     closed = GetSetProperty(W_BytesIO.closed_get_w),
+    __getstate__ = interp2app(W_BytesIO.getstate_w),
+    __setstate__ = interp2app(W_BytesIO.setstate_w),
     )