bpo-32604: Add support for a "default" arg in channel_recv(). (GH-19770)

ericsnowcurrently · web-flow · commit 5e8c691594d6 · 2020-04-28T16:11:32.000-07:00
This allows the caller to avoid creation of an exception when the channel is empty (just like `dict.get()` works). `ChannelEmptyError` is still raised if no default is provided. Automerge-Triggered-By: @ericsnowcurrently
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,19 +1350,16 @@ _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;
     }
 
     // Convert the data back to an object.
     PyObject *obj = _PyCrossInterpreterData_NewObject(data);
+    _PyCrossInterpreterData_Release(data);
+    PyMem_Free(data);
     if (obj == NULL) {
         return NULL;
     }
-    _PyCrossInterpreterData_Release(data);
-    PyMem_Free(data);
 
     return obj;
 }
@@ -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 front 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)
