core: allow Context to be pickled by non-Mitogen pickler.

mitogen/core.py

@@ -553,7 +553,7 @@ class Message(object):
         assert isinstance(self.data, BytesType)
 
     def _unpickle_context(self, context_id, name):
-        return _unpickle_context(self.router, context_id, name)
+        return _unpickle_context(context_id, name, router=self.router)
 
     def _unpickle_sender(self, context_id, dst_handle):
         return _unpickle_sender(self.router, context_id, dst_handle)
@@ -1498,14 +1498,16 @@ class Context(object):
         return 'Context(%s, %r)' % (self.context_id, self.name)
 
 
-def _unpickle_context(router, context_id, name):
-    if not (isinstance(router, Router) and
-            isinstance(context_id, (int, long)) and context_id >= 0 and (
-                (name is None) or
-                (isinstance(name, UnicodeType) and len(name) < 100))
-            ):
+def _unpickle_context(context_id, name, router=None):
+    if not (isinstance(context_id, (int, long)) and context_id >= 0 and (
+        (name is None) or
+        (isinstance(name, UnicodeType) and len(name) < 100))
+    ):
         raise TypeError('cannot unpickle Context: bad input')
-    return router.context_by_id(context_id, name=name)
+
+    if isinstance(router, Router):
+        return router.context_by_id(context_id, name=name)
+    return Context(None, context_id, name)  # For plain Jane pickle.
 
 
 class Poller(object):

tests/serialization_test.py

@@ -6,11 +6,14 @@ except ImportError:
     from StringIO import StringIO as StringIO
     from StringIO import StringIO as BytesIO
 
+import pickle
 import unittest2
 
 import mitogen.core
 from mitogen.core import b
 
+import testlib
+
 
 def roundtrip(v):
     msg = mitogen.core.Message.pickled(v)
@@ -33,5 +36,30 @@ class BlobTest(unittest2.TestCase):
         self.assertEquals(b(''), roundtrip(v))
 
 
+class ContextTest(testlib.RouterMixin, unittest2.TestCase):
+    klass = mitogen.core.Context
+
+    # Ensure Context can be round-tripped by regular pickle in addition to
+    # Mitogen's hacked pickle. Users may try to call pickle on a Context in
+    # strange circumstances, and it's often used to glue pieces of an app
+    # together (e.g. Ansible).
+
+    def test_mitogen_roundtrip(self):
+        c = self.router.fork()
+        r = mitogen.core.Receiver(self.router)
+        r.to_sender().send(c)
+        c2 = r.get().unpickle()
+        self.assertEquals(None, c2.router)
+        self.assertEquals(c.context_id, c2.context_id)
+        self.assertEquals(c.name, c2.name)
+
+    def test_vanilla_roundtrip(self):
+        c = self.router.fork()
+        c2 = pickle.loads(pickle.dumps(c))
+        self.assertEquals(None, c2.router)
+        self.assertEquals(c.context_id, c2.context_id)
+        self.assertEquals(c.name, c2.name)
+
+
 if __name__ == '__main__':
     unittest2.main()