Make interpreters.*Channel instances pickleable.

python · ericsnowcurrently · Mar 5, 2024 · Mar 1, 2024 · Mar 1, 2024 · Mar 4, 2024
commit 98a03d2a9756872182a71bf94f24cea7266dd92a
diff --git a/Lib/test/support/interpreters/channels.py b/Lib/test/support/interpreters/channels.py
@@ -38,14 +38,16 @@ class _ChannelEnd:
 
     _end = None
 
-    def __init__(self, cid):
+    def __new__(cls, cid):
+        self = super().__new__(cls)
         if self._end == 'send':
             cid = _channels._channel_id(cid, send=True, force=True)
         elif self._end == 'recv':
             cid = _channels._channel_id(cid, recv=True, force=True)
         else:
             raise NotImplementedError(self._end)
         self._id = cid
+        return self
 
     def __repr__(self):
         return f'{type(self).__name__}(id={int(self._id)})'
@@ -61,6 +63,14 @@ def __eq__(self, other):
             return NotImplemented
         return other._id == self._id
 
+    # for pickling:
+    def __getnewargs__(self):
+        return (int(self._id),)
+
+    # for pickling:
+    def __getstate__(self):
+        return None
+
     @property
     def id(self):
         return self._id

@@ -1,4 +1,5 @@
 import importlib
+import pickle
 import threading
 from textwrap import dedent
 import unittest
@@ -100,6 +101,12 @@ def test_equality(self):
         self.assertEqual(ch1, ch1)
         self.assertNotEqual(ch1, ch2)
 
+    def test_pickle(self):
+        ch, _ = channels.create()
+        data = pickle.dumps(ch)
+        unpickled = pickle.loads(data)
+        self.assertEqual(unpickled, ch)
+
 
 class TestSendChannelAttrs(TestBase):
 
@@ -125,6 +132,12 @@ def test_equality(self):
         self.assertEqual(ch1, ch1)
         self.assertNotEqual(ch1, ch2)
 
+    def test_pickle(self):
+        _, ch = channels.create()
+        data = pickle.dumps(ch)
+        unpickled = pickle.loads(data)
+        self.assertEqual(unpickled, ch)
+
 
 class TestSendRecv(TestBase):