bpo-30491: Add unawaited coroutine tracking mode

This provides a fast and easy way to get a list of coroutines that
have been created, but not awaited. For example, a test harness can
use it to reliably check after each test whether it had any missing
'await's (which can otherwise cause spurious test successes).

Doc/library/sys.rst

@@ -1292,6 +1292,66 @@ always available.
       The coroutine wrapper functionality has been deprecated, and
       will be removed in 3.8. See :issue:`32591` for details.
 
+.. function:: set_unawaited_coroutine_tracking_enabled(enabled)
+
+   Enables or disables "unawaited coroutine tracking mode". This is a
+   thread-specific value, and is initially disabled. When disabled,
+   coroutine objects are handled normally. When enabled, then every
+   coroutine is initially placed on an internal list, and the first
+   time it's iterated, it's removed from this list. At any given time,
+   therefore, the list contains all "unawaited" coroutines. For
+   example, a test suite could use this to easily and reliably check
+   after each test whether that test generated any unawaited
+   coroutines.
+
+   You can check whether the list currently contains any coroutines
+   with :func:`have_unawaited_coroutines`, and retrieve the contents
+   of the list with :func:`get_unawaited_coroutines`. Or if you want
+   to know whether tracking is currently enabled, see
+   :func:`get_unawaited_coroutine_tracking_enabled`.
+
+   .. versionadded:: 3.7
+
+   .. note::
+      This function has been added on a provisional basis (see :pep:`411`
+      for details.)  Use it only for debugging purposes.
+
+.. function:: get_unawaited_coroutine_tracking_enabled()
+
+   Checks whether unawaited coroutine tracking mode is enabled. See
+   :func:`set_unawaited_coroutine_tracking_enabled`.
+
+   .. versionadded:: 3.7
+
+   .. note::
+      This function has been added on a provisional basis (see :pep:`411`
+      for details.)  Use it only for debugging purposes.
+
+.. function:: have_unawaited_coroutines()
+
+   Returns True if unawaited coroutine tracking mode is enabled, and
+   there are currently unawaited coroutines in the internal list.
+   Otherwise, returns False. See
+   :func:`set_unawaited_coroutine_tracking_enabled`.
+
+   .. versionadded:: 3.7
+
+   .. note::
+      This function has been added on a provisional basis (see :pep:`411`
+      for details.)  Use it only for debugging purposes.
+
+.. function:: get_unawaited_coroutines()
+
+   If unawaited tracking mode is enabled, returns the internal list of
+   unawaited coroutines as a Python list object, and clears the
+   internal list. Otherwise, raises :exc:`RuntimeError`.
+
+   .. versionadded:: 3.7
+
+   .. note::
+      This function has been added on a provisional basis (see :pep:`411`
+      for details.)  Use it only for debugging purposes.
+
 .. function:: _enablelegacywindowsfsencoding()
 
    Changes the default filesystem encoding and errors mode to 'mbcs' and

Include/ceval.h

@@ -35,6 +35,10 @@ PyAPI_FUNC(void) _PyEval_SetCoroutineOriginTrackingDepth(int new_depth);
 PyAPI_FUNC(int) _PyEval_GetCoroutineOriginTrackingDepth(void);
 PyAPI_FUNC(void) _PyEval_SetCoroutineWrapper(PyObject *);
 PyAPI_FUNC(PyObject *) _PyEval_GetCoroutineWrapper(void);
+PyAPI_FUNC(void) _PyEval_SetUnawaitedCoroutineTrackingEnabled(int);
+PyAPI_FUNC(int) _PyEval_GetUnawaitedCoroutineTrackingEnabled(void);
+PyAPI_FUNC(int) _PyEval_HaveUnawaitedCoroutines(void);
+PyAPI_FUNC(PyObject *) _PyEval_GetUnawaitedCoroutines(void);
 PyAPI_FUNC(void) _PyEval_SetAsyncGenFirstiter(PyObject *);
 PyAPI_FUNC(PyObject *) _PyEval_GetAsyncGenFirstiter(void);
 PyAPI_FUNC(void) _PyEval_SetAsyncGenFinalizer(PyObject *);

Include/genobject.h

@@ -52,6 +52,14 @@ PyAPI_FUNC(void) _PyGen_Finalize(PyObject *self);
 typedef struct {
     _PyGenObject_HEAD(cr)
     PyObject *cr_origin;
+    /* Unawaited coroutine tracking. If there is another unawaited coroutine
+       after this one, cr_next_unawaited points to it. If this coroutine is in
+       the unawaited coroutine list, then cr_prev_unawaited_ptr points to
+       either the cr_next_unawaited field of the previous coroutine, or else
+       the first_unawaited_coroutine field in the relevant PyThreadState
+       object. */
+    PyObject *cr_next_unawaited;
+    PyObject **cr_prev_unawaited_ptr;
 } PyCoroObject;
 
 PyAPI_DATA(PyTypeObject) PyCoro_Type;
@@ -64,6 +72,8 @@ PyObject *_PyCoro_GetAwaitableIter(PyObject *o);
 PyAPI_FUNC(PyObject *) PyCoro_New(struct _frame *,
     PyObject *name, PyObject *qualname);
 
+PyObject * _PyCoro_PopUnawaited(void);
+
 /* Asynchronous Generators */
 
 typedef struct {

Include/pystate.h

@@ -269,6 +269,9 @@ typedef struct _ts {
     PyObject *coroutine_wrapper;
     int in_coroutine_wrapper;
 
+    int unawaited_coroutine_tracking_enabled;
+    PyObject *first_unawaited_coroutine;
+
     PyObject *async_gen_firstiter;
     PyObject *async_gen_finalizer;
 

Lib/test/test_coroutines.py

@@ -2168,6 +2168,66 @@ async def corofn():
         finally:
             warnings._warn_unawaited_coroutine = orig_wuc
 
+
+class UnawaitedTrackingTest(unittest.TestCase):
+    def test_unawaited_tracking(self):
+        was = sys.get_unawaited_coroutine_tracking_enabled()
+        try:
+            sys.set_unawaited_coroutine_tracking_enabled(False)
+            self.assertFalse(sys.get_unawaited_coroutine_tracking_enabled())
+            self.assertFalse(sys.have_unawaited_coroutines())
+            with self.assertRaises(RuntimeError):
+                sys.get_unawaited_coroutines()
+
+            sys.set_unawaited_coroutine_tracking_enabled(True)
+            self.assertTrue(sys.get_unawaited_coroutine_tracking_enabled())
+            self.assertFalse(sys.have_unawaited_coroutines())
+            self.assertEquals(sys.get_unawaited_coroutines(), [])
+
+            @types.coroutine
+            def yielder():
+                yield
+
+            async def corofn():
+                await yielder()
+
+            coro = corofn()
+            self.assertTrue(sys.have_unawaited_coroutines())
+            self.assertEquals(sys.get_unawaited_coroutines(), [coro])
+            # get consumes the coroutines
+            self.assertEquals(sys.get_unawaited_coroutines(), [])
+            coro.close()
+
+            def check_clear_fn(fn):
+                coro = corofn()
+                self.assertTrue(sys.have_unawaited_coroutines())
+                try:
+                    fn(coro)
+                except Exception:
+                    pass
+                self.assertFalse(sys.have_unawaited_coroutines())
+                coro.close()
+
+            check_clear_fn(lambda coro: coro.send(None))
+            check_clear_fn(lambda coro: coro.throw(ValueError))
+            check_clear_fn(lambda coro: coro.close())
+            check_clear_fn(lambda coro: next(coro.__await__()))
+            check_clear_fn(lambda coro: coro.__await__().send(None))
+            check_clear_fn(lambda coro: coro.__await__().throw(ValueError))
+            check_clear_fn(lambda coro: coro.__await__().close())
+
+            # Turning tracking off clears the list
+            coro = corofn()
+            self.assertTrue(sys.have_unawaited_coroutines())
+            sys.set_unawaited_coroutine_tracking_enabled(False)
+            self.assertFalse(sys.have_unawaited_coroutines())
+            sys.set_unawaited_coroutine_tracking_enabled(True)
+            self.assertFalse(sys.have_unawaited_coroutines())
+            coro.close()
+
+        finally:
+            sys.set_unawaited_coroutine_tracking_enabled(was)
+
 @support.cpython_only
 class CAPITest(unittest.TestCase):
 

Misc/NEWS.d/next/Core and Builtins/2018-01-22-23-25-47.bpo-30491.ZixJ3D.rst

@@ -0,0 +1,4 @@
+Added a new "unawaited coroutine tracking mode". For example, a test harness
+could use this to quickly and reliably check whether each test left out any
+``await``\s. See sys.{set,get}_unawaited_coroutine_enabled,
+sys.{have,get}_unawaited_coroutines.

Objects/genobject.c

@@ -147,13 +147,73 @@ gen_dealloc(PyGenObject *gen)
     PyObject_GC_Del(gen);
 }
 
+static void
+push_unawaited_coroutine(PyThreadState *tstate, PyObject *new_coro)
+{
+    Py_INCREF(new_coro);
+    PyObject *old_coro = tstate->first_unawaited_coroutine;
+    /* We have:
+         PyThreadState <-> old_coro
+
+       We want:
+         PyThreadState <-> new_coro <-> old_coro
+
+       old_coro might be NULL.
+    */
+    /* PyThreadState -> new_coro */
+    tstate->first_unawaited_coroutine = new_coro;
+    /* PyThreadState <- new_coro */
+    ((PyCoroObject *)new_coro)->cr_prev_unawaited_ptr = (
+        &tstate->first_unawaited_coroutine);
+    /* new_coro -> old_coro */
+    ((PyCoroObject *)new_coro)->cr_next_unawaited = old_coro;
+    /* new_coro <- old_coro */
+    if (old_coro) {
+        ((PyCoroObject *)old_coro)->cr_prev_unawaited_ptr = (
+            &(((PyCoroObject *)new_coro)->cr_next_unawaited));
+    }
+}
+
+static void
+gen_clear_unawaited_tracking(PyGenObject *gen)
+{
+    if (!PyCoro_CheckExact((PyObject *)gen)) {
+        return;
+    }
+    PyCoroObject *coro = (PyCoroObject *)gen;
+    if (!coro->cr_prev_unawaited_ptr) {
+        return;
+    }
+    /* We have:
+         A <-> coro <-> B
+
+       We want:
+         A <-> B
+
+       B might be NULL.
+    */
+    /* A -> B */
+    (*coro->cr_prev_unawaited_ptr) = coro->cr_next_unawaited;
+    /* A <- B */
+    if (coro->cr_next_unawaited) {
+        PyCoroObject *next = (PyCoroObject *)coro->cr_next_unawaited;
+        next->cr_prev_unawaited_ptr = coro->cr_prev_unawaited_ptr;
+    }
+    /* clear coro */
+    coro->cr_next_unawaited = NULL;
+    coro->cr_prev_unawaited_ptr = NULL;
+    Py_DECREF((PyObject *)coro);
+}
+
 static PyObject *
 gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing)
 {
     PyThreadState *tstate = PyThreadState_GET();
     PyFrameObject *f = gen->gi_frame;
     PyObject *result;
 
+    gen_clear_unawaited_tracking(gen);
+
     if (gen->gi_running) {
         const char *msg = "generator already executing";
         if (PyCoro_CheckExact(gen)) {
@@ -443,6 +503,8 @@ _gen_throw(PyGenObject *gen, int close_on_genexit,
     PyObject *yf = _PyGen_yf(gen);
     _Py_IDENTIFIER(throw);
 
+    gen_clear_unawaited_tracking(gen);
+
     if (yf) {
         PyObject *ret;
         int err;
@@ -1201,8 +1263,8 @@ PyCoro_New(PyFrameObject *f, PyObject *name, PyObject *qualname)
     }
 
     PyThreadState *tstate = PyThreadState_GET();
-    int origin_depth = tstate->coroutine_origin_tracking_depth;
 
+    int origin_depth = tstate->coroutine_origin_tracking_depth;
     if (origin_depth == 0) {
         ((PyCoroObject *)coro)->cr_origin = NULL;
     } else {
@@ -1214,6 +1276,28 @@ PyCoro_New(PyFrameObject *f, PyObject *name, PyObject *qualname)
         ((PyCoroObject *)coro)->cr_origin = cr_origin;
     }
 
+    if (tstate->unawaited_coroutine_tracking_enabled) {
+        push_unawaited_coroutine(tstate, coro);
+    } else {
+        ((PyCoroObject *)coro)->cr_next_unawaited = NULL;
+        ((PyCoroObject *)coro)->cr_prev_unawaited_ptr = NULL;
+    }
+
+    return coro;
+}
+
+/* Pops and returns an unawaited coroutine from the tracking list, or returns
+ * NULL if the list is empty. Never sets an exception.
+ */
+PyObject *
+_PyCoro_PopUnawaited(void)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    PyObject *coro = tstate->first_unawaited_coroutine;
+    if (!coro)
+        return NULL;
+    Py_INCREF(coro);
+    gen_clear_unawaited_tracking((PyGenObject *)coro);
     return coro;
 }
 

Python/ceval.c

@@ -4418,6 +4418,54 @@ _PyEval_GetCoroutineWrapper(void)
     return tstate->coroutine_wrapper;
 }
 
+void
+_PyEval_SetUnawaitedCoroutineTrackingEnabled(int new)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    if (!new) {
+        PyObject *coro;
+        while ((coro = _PyCoro_PopUnawaited()))
+            Py_DECREF(coro);
+    }
+    tstate->unawaited_coroutine_tracking_enabled = new;
+}
+
+int
+_PyEval_GetUnawaitedCoroutineTrackingEnabled(void)
+{
+    return PyThreadState_GET()->unawaited_coroutine_tracking_enabled;
+}
+
+int
+_PyEval_HaveUnawaitedCoroutines(void)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    return (tstate->first_unawaited_coroutine != NULL);
+}
+
+PyObject *
+_PyEval_GetUnawaitedCoroutines(void)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    if (!tstate->unawaited_coroutine_tracking_enabled) {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "unawaited coroutine tracking not enabled");
+        return NULL;
+    }
+    PyObject *output = PyList_New(0);
+    if (!output)
+        return NULL;
+    PyObject *coro;
+    while ((coro = _PyCoro_PopUnawaited())) {
+        if (PyList_Append(output, coro) < 0) {
+            Py_DECREF(coro);
+            return NULL;
+        }
+        Py_DECREF(coro);
+    }
+    return output;
+}
+
 void
 _PyEval_SetAsyncGenFirstiter(PyObject *firstiter)
 {