issue #218: core: add Secret and Blob types.

docs/getting_started.rst

@@ -320,6 +320,8 @@ remote procedure calls:
 
 User-defined types may not be used, except for:
 
+* :py:class:`mitogen.core.Blob`
+* :py:class:`mitogen.core.Secret`
 * :py:class:`mitogen.core.CallError`
 * :py:class:`mitogen.core.Context`
 * :py:class:`mitogen.core.Sender`

mitogen/core.py

@@ -81,8 +81,12 @@ IS_DEAD = 999
 PY3 = sys.version_info > (3,)
 if PY3:
     b = lambda s: s.encode('latin-1')
+    BytesType = bytes
+    UnicodeType = unicode
 else:
     b = str
+    BytesType = str
+    UnicodeType = unicode
 
 CHUNK_SIZE = 131072
 _tls = threading.local()
@@ -109,6 +113,25 @@ class LatchError(Error):
     pass
 
 
+class Blob(BytesType):
+    def __repr__(self):
+        return '[blob: %d bytes]' % len(self)
+
+    def __reduce__(self):
+        return (_unpickle_blob, (BytesType(self),))
+
+
+class Secret(UnicodeType):
+    def __repr__(self):
+        return '[secret]'
+
+    def __str__(self):
+        return UnicodeType(self)
+
+    def __reduce__(self):
+        return (_unpickle_secret, (UnicodeType(self),))
+
+
 class CallError(Error):
     def __init__(self, fmt=None, *args):
         if not isinstance(fmt, Exception):
@@ -127,6 +150,14 @@ class CallError(Error):
         return (_unpickle_call_error, (self[0],))
 
 
+def _unpickle_blob(s):
+    return Blob(s)
+
+
+def _unpickle_secret(s):
+    return Secret(s)
+
+
 def _unpickle_call_error(s):
     if not (type(s) is str and len(s) < 10000):
         raise TypeError('cannot unpickle CallError: bad input')
@@ -310,7 +341,10 @@ class Message(object):
                 return self._unpickle_sender
             elif func == '_unpickle_context':
                 return self._unpickle_context
-
+            elif func == '_unpickle_blob':
+                return _unpickle_blob
+            elif func == '_unpickle_secret':
+                return _unpickle_secret
         raise StreamError('cannot unpickle %r/%r', module, func)
 
     @property

mitogen/utils.py

@@ -101,19 +101,25 @@ def with_router(func):
     return wrapper
 
 
+PASSTHROUGH = (
+    int, float, bool,
+    type(None),
+    mitogen.core.Context,
+    mitogen.core.CallError,
+    mitogen.core.Blob,
+    mitogen.core.Secret,
+)
+
 def cast(obj):
     if isinstance(obj, dict):
         return {cast(k): cast(v) for k, v in obj.iteritems()}
     if isinstance(obj, (list, tuple)):
         return [cast(v) for v in obj]
-    if obj is None or isinstance(obj, (int, float)):
+    if isinstance(obj, PASSTHROUGH):
         return obj
     if isinstance(obj, unicode):
         return unicode(obj)
     if isinstance(obj, str):
         return str(obj)
-    if isinstance(obj, (mitogen.core.Context,
-                        mitogen.core.CallError)):
-        return obj
 
     raise TypeError("Cannot serialize: %r: %r" % (type(obj), obj))

tests/types_test.py

@@ -0,0 +1,72 @@
+
+import cStringIO
+
+import unittest2
+
+import mitogen.core
+
+
+class BlobTest(unittest2.TestCase):
+    klass = mitogen.core.Blob
+
+    def make(self):
+        return self.klass('x' * 128)
+
+    def test_repr(self):
+        blob = self.make()
+        self.assertEquals('[blob: 128 bytes]', repr(blob))
+
+    def test_decays_on_constructor(self):
+        blob = self.make()
+        self.assertEquals('x'*128, mitogen.core.BytesType(blob))
+
+    def test_decays_on_write(self):
+        blob = self.make()
+        io = cStringIO.StringIO()
+        io.write(blob)
+        self.assertEquals(128, io.tell())
+        self.assertEquals('x'*128, io.getvalue())
+
+    def test_message_roundtrip(self):
+        blob = self.make()
+        msg = mitogen.core.Message.pickled(blob)
+        blob2 = msg.unpickle()
+        self.assertEquals(type(blob), type(blob2))
+        self.assertEquals(repr(blob), repr(blob2))
+        self.assertEquals(mitogen.core.BytesType(blob),
+                          mitogen.core.BytesType(blob2))
+
+
+class SecretTest(unittest2.TestCase):
+    klass = mitogen.core.Secret
+
+    def make(self):
+        return self.klass('password')
+
+    def test_repr(self):
+        secret = self.make()
+        self.assertEquals('[secret]', repr(secret))
+
+    def test_decays_on_constructor(self):
+        secret = self.make()
+        self.assertEquals('password', mitogen.core.UnicodeType(secret))
+
+    def test_decays_on_write(self):
+        secret = self.make()
+        io = cStringIO.StringIO()
+        io.write(secret)
+        self.assertEquals(8, io.tell())
+        self.assertEquals('password', io.getvalue())
+
+    def test_message_roundtrip(self):
+        secret = self.make()
+        msg = mitogen.core.Message.pickled(secret)
+        secret2 = msg.unpickle()
+        self.assertEquals(type(secret), type(secret2))
+        self.assertEquals(repr(secret), repr(secret2))
+        self.assertEquals(mitogen.core.BytesType(secret),
+                          mitogen.core.BytesType(secret2))
+
+
+if __name__ == '__main__':
+    unittest2.main()