Add support for a default arg in channel_recv().

python · miss-islington · Apr 28, 2020 · Apr 28, 2020 · Apr 28, 2020 · Apr 28, 2020
commit b5edb1211a9e5096257a0c428c9ded4ef2b7169b
diff --git a/Lib/test/test__xxsubinterpreters.py b/Lib/test/test__xxsubinterpreters.py
@@ -1302,6 +1302,27 @@ def test_recv_empty(self):
         with self.assertRaises(interpreters.ChannelEmptyError):
             interpreters.channel_recv(cid)
 
+    def test_recv_default(self):
+        default = object()
+        cid = interpreters.channel_create()
+        obj1 = interpreters.channel_recv(cid, default)
+        interpreters.channel_send(cid, None)
+        interpreters.channel_send(cid, 1)
+        interpreters.channel_send(cid, b'spam')
+        interpreters.channel_send(cid, b'eggs')
+        obj2 = interpreters.channel_recv(cid, default)
+        obj3 = interpreters.channel_recv(cid, default)
+        obj4 = interpreters.channel_recv(cid)
+        obj5 = interpreters.channel_recv(cid, default)
+        obj6 = interpreters.channel_recv(cid, default)
+
+        self.assertIs(obj1, default)
+        self.assertIs(obj2, None)
+        self.assertEqual(obj3, 1)
+        self.assertEqual(obj4, b'spam')
+        self.assertEqual(obj5, b'eggs')
+        self.assertIs(obj6, default)
+
     def test_run_string_arg_unresolved(self):
         cid = interpreters.channel_create()
         interp = interpreters.create()

diff --git a/Modules/_xxsubinterpretersmodule.c b/Modules/_xxsubinterpretersmodule.c
@@ -1350,9 +1350,6 @@ _channel_recv(_channels *channels, int64_t id)
     _PyCrossInterpreterData *data = _channel_next(chan, PyInterpreterState_GetID(interp));
     PyThread_release_lock(mutex);
     if (data == NULL) {
-        if (!PyErr_Occurred()) {
-            PyErr_Format(ChannelEmptyError, "channel %" PRId64 " is empty", id);
-        }
         return NULL;
     }
 
@@ -2351,20 +2348,37 @@ Add the object's data to the channel's queue.");
 static PyObject *
 channel_recv(PyObject *self, PyObject *args, PyObject *kwds)
 {
-    static char *kwlist[] = {"cid", NULL};
+    static char *kwlist[] = {"cid", "default", NULL};
     int64_t cid;
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&:channel_recv", kwlist,
-                                     channel_id_converter, &cid)) {
+    PyObject *dflt = NULL;
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&|O:channel_recv", kwlist,
+                                     channel_id_converter, &cid, &dflt)) {
         return NULL;
     }
+    Py_XINCREF(dflt);
 
-    return _channel_recv(&_globals.channels, cid);
+    PyObject *obj = _channel_recv(&_globals.channels, cid);
+    if (obj != NULL) {
+        Py_XDECREF(dflt);
+        return obj;
+    } else if (PyErr_Occurred()) {
+        Py_XDECREF(dflt);
+        return NULL;
+    } else if (dflt != NULL) {
+        return dflt;
+    } else {
+        PyErr_Format(ChannelEmptyError, "channel %" PRId64 " is empty", cid);
+        return NULL;
+    }
 }
 
 PyDoc_STRVAR(channel_recv_doc,
-"channel_recv(cid) -> obj\n\
+"channel_recv(cid, [default]) -> obj\n\
+\n\
+Return a new object from the data at the from of the channel's queue.\n\
 \n\
-Return a new object from the data at the from of the channel's queue.");
+If there is nothing to receive then raise ChannelEmptyError, unless\n\
+a default value is provided.  In that case return it.");
 
 static PyObject *
 channel_close(PyObject *self, PyObject *args, PyObject *kwds)