Call _channels._channel_id() in _ChannelEnd.__init__().

ericsnowcurrently · Oct 2, 2023 · Sep 11, 2023 · Sep 11, 2023 · Sep 12, 2023 · Sep 12, 2023
commit e424ab556ef40ed12007fbe158e76077f916e404
diff --git a/Lib/test/support/interpreters.py b/Lib/test/support/interpreters.py
@@ -118,10 +118,16 @@ def list_all_channels():
 class _ChannelEnd:
     """The base class for RecvChannel and SendChannel."""
 
-    def __init__(self, id):
-        if not isinstance(id, (int, _channels.ChannelID)):
-            raise TypeError(f'id must be an int, got {id!r}')
-        self._id = id
+    _end = None
+
+    def __init__(self, cid):
+        if self._end == 'send':
+            cid = _channels._channel_id(cid, send=True)
+        elif self._end == 'recv':
+            cid = _channels._channel_id(cid, recv=True)
+        else:
+            raise NotImplementedError(self._end)
+        self._id = cid
 
     def __repr__(self):
         return f'{type(self).__name__}(id={int(self._id)})'
@@ -148,6 +154,8 @@ def id(self):
 class RecvChannel(_ChannelEnd):
     """The receiving end of a cross-interpreter channel."""
 
+    _end = 'recv'
+
     def recv(self, *, _sentinel=object(), _delay=10 / 1000):  # 10 milliseconds
         """Return the next object from the channel.
 
@@ -179,6 +187,8 @@ def close(self):
 class SendChannel(_ChannelEnd):
     """The sending end of a cross-interpreter channel."""
 
+    _end = 'send'
+
     def send(self, obj):
         """Send the object (i.e. its data) to the channel's receiving end.
 
@@ -203,3 +213,6 @@ def send_nowait(self, obj):
 
     def close(self):
         _channels.close(self._id, send=True)
+
+
+_channels._register_end_types(SendChannel, RecvChannel)
diff --git a/Modules/_xxinterpchannelsmodule.c b/Modules/_xxinterpchannelsmodule.c
@@ -1529,17 +1529,20 @@ typedef struct channelid {
 struct channel_id_converter_data {
     PyObject *module;
     int64_t cid;
+    int end;
 };
 
 static int
 channel_id_converter(PyObject *arg, void *ptr)
 {
     int64_t cid;
+    int end = 0;
     struct channel_id_converter_data *data = ptr;
     module_state *state = get_module_state(data->module);
     assert(state != NULL);
     if (PyObject_TypeCheck(arg, state->ChannelIDType)) {
         cid = ((channelid *)arg)->id;
+        end = ((channelid *)arg)->end;
     }
     else if (PyIndex_Check(arg)) {
         cid = PyLong_AsLongLong(arg);
@@ -1559,6 +1562,7 @@ channel_id_converter(PyObject *arg, void *ptr)
         return 0;
     }
     data->cid = cid;
+    data->end = end;
     return 1;
 }
 
@@ -1600,6 +1604,7 @@ _channelid_new(PyObject *mod, PyTypeObject *cls,
 {
     static char *kwlist[] = {"id", "send", "recv", "force", "_resolve", NULL};
     int64_t cid;
+    int end;
     struct channel_id_converter_data cid_data = {
         .module = mod,
     };
@@ -1614,21 +1619,25 @@ _channelid_new(PyObject *mod, PyTypeObject *cls,
         return NULL;
     }
     cid = cid_data.cid;
+    end = cid_data.end;
 
     // Handle "send" and "recv".
     if (send == 0 && recv == 0) {
         PyErr_SetString(PyExc_ValueError,
                         "'send' and 'recv' cannot both be False");
         return NULL;
     }
-
-    int end = 0;
-    if (send == 1) {
+    else if (send == 1) {
         if (recv == 0 || recv == -1) {
             end = CHANNEL_SEND;
         }
+        else {
+            assert(recv == 1);
+            end = 0;
+        }
     }
     else if (recv == 1) {
+        assert(send == 0 || send == -1);
         end = CHANNEL_RECV;
     }
 
@@ -2346,13 +2355,9 @@ channel__channel_id(PyObject *self, PyObject *args, PyObject *kwds)
         return NULL;
     }
     PyTypeObject *cls = state->ChannelIDType;
-    PyObject *mod = get_module_from_owned_type(cls);
-    if (mod == NULL) {
-        return NULL;
-    }
-    PyObject *cid = _channelid_new(mod, cls, args, kwds);
-    Py_DECREF(mod);
-    return cid;
+    assert(get_module_from_owned_type(cls) == self);
+
+    return _channelid_new(self, cls, args, kwds);
 }
 
 static PyMethodDef module_functions[] = {