diff --git a/mitogen/core.py b/mitogen/core.py index 593b9d21..51350351 100644 --- a/mitogen/core.py +++ b/mitogen/core.py @@ -425,6 +425,8 @@ class Message(object): return Kwargs elif module == '_codecs' and func == 'encode': return self._unpickle_bytes + elif module == '__builtin__' and func == 'bytes': + return BytesType raise StreamError('cannot unpickle %r/%r', module, func) @property diff --git a/tests/serialization_test.py b/tests/serialization_test.py new file mode 100644 index 00000000..3b56e10a --- /dev/null +++ b/tests/serialization_test.py @@ -0,0 +1,37 @@ + +try: + from io import StringIO + from io import BytesIO +except ImportError: + from StringIO import StringIO as StringIO + from StringIO import StringIO as BytesIO + +import unittest2 + +import mitogen.core +from mitogen.core import b + + +def roundtrip(v): + msg = mitogen.core.Message.pickled(v) + return mitogen.core.Message(data=msg.data).unpickle() + + +class BlobTest(unittest2.TestCase): + klass = mitogen.core.Blob + + # Python 3 pickle protocol 2 does weird stuff depending on whether an empty + # or nonempty bytes is being serialized. For non-empty, it yields a + # _codecs.encode() call. For empty, it yields a bytes() call. + + def test_nonempty_bytes(self): + v = mitogen.core.Blob(b('dave')) + self.assertEquals(b('dave'), roundtrip(v)) + + def test_empty_bytes(self): + v = mitogen.core.Blob(b('')) + self.assertEquals(b(''), roundtrip(v)) + + +if __name__ == '__main__': + unittest2.main()