From 0f29baa0777671b5d1bff28a0688dbb882416629 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Thu, 29 Mar 2018 15:32:01 +0545
Subject: [PATCH] core: support pickling senders, Receiver.to_sender()

CC @moreati, in case this impacts you
---
 mitogen/core.py             | 22 +++++++++++++++++++++-
 tests/call_function_test.py | 19 ++++++++++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/mitogen/core.py b/mitogen/core.py
index f6ddf23d..5739e2b4 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -299,6 +299,9 @@ class Message(object):
     def _unpickle_context(self, context_id, name):
         return _unpickle_context(self.router, context_id, name)
 
+    def _unpickle_sender(self, context_id, dst_handle):
+        return _unpickle_sender(self.router, context_id, dst_handle)
+
     def _find_global(self, module, func):
         """Return the class implementing `module_name.class_name` or raise
         `StreamError` if the module is not whitelisted."""
@@ -307,6 +310,8 @@ class Message(object):
                 return _unpickle_call_error
             elif func == '_unpickle_dead':
                 return _unpickle_dead
+            elif func == '_unpickle_sender':
+                return self._unpickle_sender
             elif func == '_unpickle_context':
                 return self._unpickle_context
 
@@ -366,6 +371,9 @@ class Sender(object):
     def __repr__(self):
         return 'Sender(%r, %r)' % (self.context, self.dst_handle)
 
+    def __reduce__(self):
+        return _unpickle_sender, (self.context.context_id, self.dst_handle)
+
     def close(self):
         """Indicate this channel is closed to the remote side."""
         _vv and IOLOG.debug('%r.close()', self)
@@ -387,6 +395,14 @@ class Sender(object):
         )
 
 
+def _unpickle_sender(router, context_id, dst_handle):
+    if not (isinstance(router, Router) and
+            isinstance(context_id, (int, long)) and context_id >= 0 and
+            isinstance(dst_handle, (int, long)) and dst_handle > 0):
+        raise TypeError('cannot unpickle Sender: bad input')
+    return Sender(Context(router, context_id), dst_handle)
+
+
 class Receiver(object):
     notify = None
     raise_channelerror = True
@@ -401,6 +417,10 @@ class Receiver(object):
     def __repr__(self):
         return 'Receiver(%r, %r)' % (self.router, self.handle)
 
+    def to_sender(self):
+        context = Context(self.router, mitogen.context_id)
+        return Sender(context, self.handle)
+
     def _on_receive(self, msg):
         """Callback from the Stream; appends data to the internal queue."""
         _vv and IOLOG.debug('%r._on_receive(%r)', self, msg)
@@ -912,7 +932,7 @@ class Context(object):
 
 def _unpickle_context(router, context_id, name):
     if not (isinstance(router, Router) and
-            isinstance(context_id, (int, long)) and context_id > 0 and
+            isinstance(context_id, (int, long)) and context_id >= 0 and
             isinstance(name, basestring) and len(name) < 100):
         raise TypeError('cannot unpickle Context: bad input')
     return router.context_class(router, context_id, name)
diff --git a/tests/call_function_test.py b/tests/call_function_test.py
index cc6c32b6..301b7798 100644
--- a/tests/call_function_test.py
+++ b/tests/call_function_test.py
@@ -33,10 +33,16 @@ def func_accepts_returns_context(context):
     return context
 
 
+def func_accepts_returns_sender(sender):
+    sender.put(123)
+    sender.close()
+    return sender
+
+
 class CallFunctionTest(testlib.RouterMixin, testlib.TestCase):
     def setUp(self):
         super(CallFunctionTest, self).setUp()
-        self.local = self.router.local()
+        self.local = self.router.fork()
 
     def test_succeeds(self):
         self.assertEqual(3, self.local.call(function_that_adds_numbers, 1, 2))
@@ -87,6 +93,17 @@ class CallFunctionTest(testlib.RouterMixin, testlib.TestCase):
         self.assertEqual(context.context_id, self.local.context_id)
         self.assertEqual(context.name, self.local.name)
 
+    def test_accepts_returns_sender(self):
+        recv = mitogen.core.Receiver(self.router)
+        sender = recv.to_sender()
+        sender2 = self.local.call(func_accepts_returns_sender, sender)
+        self.assertEquals(sender.context.context_id,
+                          sender2.context.context_id)
+        self.assertEquals(sender.dst_handle, sender2.dst_handle)
+        self.assertEquals(123, recv.get().unpickle())
+        self.assertRaises(mitogen.core.ChannelError,
+                          lambda: recv.get().unpickle())
+
 
 if __name__ == '__main__':
     unittest2.main()