Add coro.cr_origin and sys.set_coroutine_origin_tracking_depth

python · 1st1 · Jan 21, 2018 · Jan 20, 2018 · Jan 20, 2018 · Jan 20, 2018
commit 091dc24c365786131bbe056287a7dae1219fccd5
diff --git a/Doc/library/inspect.rst b/Doc/library/inspect.rst
@@ -215,6 +215,10 @@ attributes:
 +-----------+-------------------+---------------------------+
 |           | cr_code           | code                      |
 +-----------+-------------------+---------------------------+
+|           | cr_origin         | where coroutine was       |
+|           |                   | created, if coroutine     |
+|           |                   | origin tracking is enabled|
++-----------+-------------------+---------------------------+
 | builtin   | __doc__           | documentation string      |
 +-----------+-------------------+---------------------------+
 |           | __name__          | original name of this     |
@@ -234,6 +238,9 @@ attributes:
    The ``__name__`` attribute of generators is now set from the function
    name, instead of the code name, and it can now be modified.
 
+.. versionchanged:: 3.7
+
+   Add ``cr_origin`` attribute to coroutines.
 
 .. function:: getmembers(object[, predicate])
 

diff --git a/Doc/library/sys.rst b/Doc/library/sys.rst
@@ -1212,6 +1212,27 @@ always available.
       This function has been added on a provisional basis (see :pep:`411`
       for details.)
 
+.. function:: set_coroutine_origin_tracking_depth(depth)
+
+   Allows enabling or disabling coroutine origin tracking. When
+   enabled, the ``cr_origin`` attribute on coroutine objects will
+   contain a list of (filename, line number, function name) tuples
+   describing the traceback where the coroutine object was created.
+   When disabled, ``cr_origin`` will be None.
+
+   To enable, pass a *depth* value greater than zero; this sets the
+   number of frames whose information will be captured. To disable,
+   pass set *depth* to zero.
+
+   Returns the old value of *depth*.
+
+   This setting is thread-local.
+
+   .. 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:: set_coroutine_wrapper(wrapper)
 

diff --git a/Include/ceval.h b/Include/ceval.h
@@ -31,6 +31,7 @@ PyAPI_FUNC(PyObject *) PyEval_CallMethod(PyObject *obj,
 #ifndef Py_LIMITED_API
 PyAPI_FUNC(void) PyEval_SetProfile(Py_tracefunc, PyObject *);
 PyAPI_FUNC(void) PyEval_SetTrace(Py_tracefunc, PyObject *);
+PyAPI_FUNC(int) _PyEval_SetCoroutineOriginTrackingDepth(int new_depth);
 PyAPI_FUNC(void) _PyEval_SetCoroutineWrapper(PyObject *);
 PyAPI_FUNC(PyObject *) _PyEval_GetCoroutineWrapper(void);
 PyAPI_FUNC(void) _PyEval_SetAsyncGenFirstiter(PyObject *);

diff --git a/Include/genobject.h b/Include/genobject.h
@@ -51,6 +51,7 @@ PyAPI_FUNC(void) _PyGen_Finalize(PyObject *self);
 #ifndef Py_LIMITED_API
 typedef struct {
     _PyGenObject_HEAD(cr)
+    PyObject *cr_origin;
 } PyCoroObject;
 
 PyAPI_DATA(PyTypeObject) PyCoro_Type;

diff --git a/Include/pystate.h b/Include/pystate.h
@@ -262,6 +262,8 @@ typedef struct _ts {
     void (*on_delete)(void *);
     void *on_delete_data;
 
+    int coroutine_origin_tracking_depth;
+
     PyObject *coroutine_wrapper;
     int in_coroutine_wrapper;
 

diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py
@@ -1,4 +1,4 @@
-import contextlib
+from contextlib import contextmanager, closing
 import copy
 import inspect
 import pickle
@@ -58,7 +58,7 @@ def run_async__await__(coro):
     return buffer, result
 
 
-@contextlib.contextmanager
+@contextmanager
 def silence_coro_gc():
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
@@ -2041,6 +2041,51 @@ def wrap(gen):
             sys.set_coroutine_wrapper(None)
 
 
+class OriginTrackingTest(unittest.TestCase):
+    def here(self):
+        info = inspect.getframeinfo(inspect.currentframe().f_back)
+        return (info.filename, info.lineno)
+
+    def test_origin_tracking(self):
+        orig_depth = sys.set_coroutine_origin_tracking_depth(0)
+        try:
+            async def corofn():
+                pass
+
+            with closing(corofn()) as coro:
+                self.assertIsNone(coro.cr_origin)
+
+            self.assertEqual(sys.set_coroutine_origin_tracking_depth(1), 0)
+
+            fname, lineno = self.here()
+            with closing(corofn()) as coro:
+                self.assertEqual(coro.cr_origin,
+                                 [(fname, lineno + 1, "test_origin_tracking")])
+
+            self.assertEqual(sys.set_coroutine_origin_tracking_depth(2), 1)
+            def nested():
+                return (self.here(), corofn())
+            fname, lineno = self.here()
+            ((nested_fname, nested_lineno), coro) = nested()
+            with closing(coro):
+                self.assertEqual(coro.cr_origin,
+                                 [(nested_fname, nested_lineno, "nested"),
+                                  (fname, lineno + 1, "test_origin_tracking")])
+
+            # Check we handle running out of frames correctly
+            sys.set_coroutine_origin_tracking_depth(1000)
+            with closing(corofn()) as coro:
+                self.assertTrue(2 < len(coro.cr_origin) < 1000)
+
+            # We can't set depth negative
+            with self.assertRaises(ValueError):
+                sys.set_coroutine_origin_tracking_depth(-1)
+            # And trying leaves it unchanged
+            self.assertEqual(sys.set_coroutine_origin_tracking_depth(0), 1000)
+
+        finally:
+            sys.set_coroutine_origin_tracking_depth(orig_depth)
+
 @support.cpython_only
 class CAPITest(unittest.TestCase):
 

diff --git a/Objects/genobject.c b/Objects/genobject.c
@@ -32,6 +32,8 @@ gen_traverse(PyGenObject *gen, visitproc visit, void *arg)
     Py_VISIT(gen->gi_code);
     Py_VISIT(gen->gi_name);
     Py_VISIT(gen->gi_qualname);
+    if (((PyCodeObject *)gen->gi_code)->co_flags & CO_COROUTINE)
+        Py_VISIT(((PyCoroObject *)gen)->cr_origin);
     return exc_state_traverse(&gen->gi_exc_state, visit, arg);
 }
 
@@ -137,6 +139,8 @@ gen_dealloc(PyGenObject *gen)
         gen->gi_frame->f_gen = NULL;
         Py_CLEAR(gen->gi_frame);
     }
+    if (((PyCodeObject *)gen->gi_code)->co_flags & CO_COROUTINE)
+        Py_CLEAR(((PyCoroObject *)gen)->cr_origin);
     Py_CLEAR(gen->gi_code);
     Py_CLEAR(gen->gi_name);
     Py_CLEAR(gen->gi_qualname);
@@ -990,6 +994,7 @@ static PyMemberDef coro_memberlist[] = {
     {"cr_frame",     T_OBJECT, offsetof(PyCoroObject, cr_frame),    READONLY},
     {"cr_running",   T_BOOL,   offsetof(PyCoroObject, cr_running),  READONLY},
     {"cr_code",      T_OBJECT, offsetof(PyCoroObject, cr_code),     READONLY},
+    {"cr_origin",    T_OBJECT, offsetof(PyCoroObject, cr_origin),   READONLY},
     {NULL}      /* Sentinel */
 };
 
@@ -1161,7 +1166,47 @@ PyTypeObject _PyCoroWrapper_Type = {
 PyObject *
 PyCoro_New(PyFrameObject *f, PyObject *name, PyObject *qualname)
 {
-    return gen_new_with_qualname(&PyCoro_Type, f, name, qualname);
+    PyObject *coro = gen_new_with_qualname(&PyCoro_Type, f, name, qualname);
+    if (!coro)
+        return NULL;
+
+    PyThreadState *tstate = PyThreadState_GET();
+    int depth = tstate->coroutine_origin_tracking_depth;
+
+    if (depth == 0) {
+        ((PyCoroObject *)coro)->cr_origin = NULL;
+    } else {
+        PyObject *origin = PyList_New(depth);
+        /* Immediately pass ownership to coro, so on error paths we don't have
+           to worry about it separately. */
+        ((PyCoroObject *)coro)->cr_origin = origin;
+        PyFrameObject *frame = PyEval_GetFrame();
+        int i = 0;
+        for (; i < depth; ++i) {
+            if (!frame)
+                break;
+
+            PyObject *frameinfo = Py_BuildValue(
+                "OiO",
+                frame->f_code->co_filename,
+                PyFrame_GetLineNumber(frame),
+                frame->f_code->co_name);
+            if (!frameinfo) {
+                Py_DECREF(coro);
+                return NULL;
+            }
+            PyList_SET_ITEM(origin, i, frameinfo);
+
+            frame = frame->f_back;
+        }
+        /* Truncate the list if necessary */
+        if (PyList_SetSlice(origin, i, depth, NULL) < 0) {
+            Py_DECREF(coro);
+            return NULL;
+        }
+    }
+
+    return coro;
 }
 
 

diff --git a/Python/ceval.c b/Python/ceval.c
@@ -4387,6 +4387,15 @@ PyEval_SetTrace(Py_tracefunc func, PyObject *arg)
                            || (tstate->c_profilefunc != NULL));
 }
 
+int
+_PyEval_SetCoroutineOriginTrackingDepth(int new_depth)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    int old_depth = tstate->coroutine_origin_tracking_depth;
+    tstate->coroutine_origin_tracking_depth = new_depth;
+    return old_depth;
+}
+
 void
 _PyEval_SetCoroutineWrapper(PyObject *wrapper)
 {

diff --git a/Python/clinic/sysmodule.c.h b/Python/clinic/sysmodule.c.h
@@ -0,0 +1,43 @@
+/*[clinic input]
+preserve
+[clinic start generated code]*/
+
+PyDoc_STRVAR(sys_set_coroutine_origin_tracking_depth__doc__,
+"set_coroutine_origin_tracking_depth($module, /, depth)\n"
+"--\n"
+"\n"
+"Enable or disable origin tracking for coroutine objects in this thread.\n"
+"\n"
+"Coroutine objects will track \'depth\' frames of traceback information about\n"
+"where they came from, available in their cr_origin attribute. Set depth of 0\n"
+"to disable. Returns old value.");
+
+#define SYS_SET_COROUTINE_ORIGIN_TRACKING_DEPTH_METHODDEF    \
+    {"set_coroutine_origin_tracking_depth", (PyCFunction)sys_set_coroutine_origin_tracking_depth, METH_FASTCALL|METH_KEYWORDS, sys_set_coroutine_origin_tracking_depth__doc__},
+
+static int
+sys_set_coroutine_origin_tracking_depth_impl(PyObject *module, int depth);
+
+static PyObject *
+sys_set_coroutine_origin_tracking_depth(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    PyObject *return_value = NULL;
+    static const char * const _keywords[] = {"depth", NULL};
+    static _PyArg_Parser _parser = {"i:set_coroutine_origin_tracking_depth", _keywords, 0};
+    int depth;
+    int _return_value;
+
+    if (!_PyArg_ParseStackAndKeywords(args, nargs, kwnames, &_parser,
+        &depth)) {
+        goto exit;
+    }
+    _return_value = sys_set_coroutine_origin_tracking_depth_impl(module, depth);
+    if ((_return_value == -1) && PyErr_Occurred()) {
+        goto exit;
+    }
+    return_value = PyLong_FromLong((long)_return_value);
+
+exit:
+    return return_value;
+}
+/*[clinic end generated code: output=9448b0765e4d5bf0 input=a9049054013a1b77]*/
diff --git a/Python/pystate.c b/Python/pystate.c
@@ -305,6 +305,8 @@ new_threadstate(PyInterpreterState *interp, int init)
         tstate->on_delete = NULL;
         tstate->on_delete_data = NULL;
 
+        tstate->coroutine_origin_tracking_depth = 0;
+
         tstate->coroutine_wrapper = NULL;
         tstate->in_coroutine_wrapper = 0;
 

diff --git a/Python/sysmodule.c b/Python/sysmodule.c
@@ -34,6 +34,13 @@ extern void *PyWin_DLLhModule;
 extern const char *PyWin_DLLVersionString;
 #endif
 
+/*[clinic input]
+module sys
+[clinic start generated code]*/
+/*[clinic end generated code: output=da39a3ee5e6b4b0d input=3726b388feee8cea]*/
+
+#include "clinic/sysmodule.c.h"
+
 _Py_IDENTIFIER(_);
 _Py_IDENTIFIER(__sizeof__);
 _Py_IDENTIFIER(_xoptions);
@@ -710,6 +717,29 @@ sys_setrecursionlimit(PyObject *self, PyObject *args)
     Py_RETURN_NONE;
 }
 
+/*[clinic input]
+sys.set_coroutine_origin_tracking_depth -> int
+
+  depth: int
+
+Enable or disable origin tracking for coroutine objects in this thread.
+
+Coroutine objects will track 'depth' frames of traceback information about
+where they came from, available in their cr_origin attribute. Set depth of 0
+to disable. Returns old value.
+[clinic start generated code]*/
+
+static int
+sys_set_coroutine_origin_tracking_depth_impl(PyObject *module, int depth)
+/*[clinic end generated code: output=1d421106d83a2fce input=f0a48609fc463a5c]*/
+{
+    if (depth < 0) {
+        PyErr_SetString(PyExc_ValueError, "depth must be >= 0");
+        return -1;
+    }
+    return _PyEval_SetCoroutineOriginTrackingDepth(depth);
+}
+
 static PyObject *
 sys_set_coroutine_wrapper(PyObject *self, PyObject *wrapper)
 {
@@ -1512,6 +1542,7 @@ static PyMethodDef sys_methods[] = {
     {"call_tracing", sys_call_tracing, METH_VARARGS, call_tracing_doc},
     {"_debugmallocstats", sys_debugmallocstats, METH_NOARGS,
      debugmallocstats_doc},
+    SYS_SET_COROUTINE_ORIGIN_TRACKING_DEPTH_METHODDEF
     {"set_coroutine_wrapper", sys_set_coroutine_wrapper, METH_O,
      set_coroutine_wrapper_doc},
     {"get_coroutine_wrapper", sys_get_coroutine_wrapper, METH_NOARGS,