[Python-checkins] gh-100344: Add C implementation for `asyncio.current_task` (#100345)

Thu Dec 22 08:38:19 EST 2022

https://github.com/python/cpython/commit/4cc63e0d4e4cf3299dcc0ea81616ba072ae5589d
commit: 4cc63e0d4e4cf3299dcc0ea81616ba072ae5589d
branch: main
author: Itamar Ostricher <itamarost at gmail.com>
committer: kumaraditya303 <59607654+kumaraditya303 at users.noreply.github.com>
date: 2022年12月22日T19:08:12+05:30
summary:
gh-100344: Add C implementation for `asyncio.current_task` (#100345)
Co-authored-by: pranavtbhat
files:
A Misc/NEWS.d/next/Library/2022-12-19-12-18-28.gh-issue-100344.lfCqpE.rst
M Doc/whatsnew/3.12.rst
M Lib/asyncio/tasks.py
M Lib/test/test_asyncio/test_tasks.py
M Modules/_asynciomodule.c
M Modules/clinic/_asynciomodule.c.h

diff --git a/Doc/whatsnew/3.12.rst b/Doc/whatsnew/3.12.rst
index 0cc4471364b6..d480be27542c 100644
--- a/Doc/whatsnew/3.12.rst
+++ b/Doc/whatsnew/3.12.rst
@@ -225,6 +225,9 @@ asyncio
 a custom event loop factory.
 (Contributed by Kumar Aditya in :gh:`99388`.)
 
+* Add C implementation of :func:`asyncio.current_task` for 4x-6x speedup.
+ (Contributed by Itamar Ostricher and Pranav Thulasiram Bhat in :gh:`100344`.)
+
 inspect
 -------
 
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index fa853283c0c5..e78719de216f 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -964,6 +964,7 @@ def _unregister_task(task):
 _all_tasks.discard(task)
 
 
+_py_current_task = current_task
 _py_register_task = _register_task
 _py_unregister_task = _unregister_task
 _py_enter_task = _enter_task
@@ -973,10 +974,12 @@ def _unregister_task(task):
 try:
 from _asyncio import (_register_task, _unregister_task,
 _enter_task, _leave_task,
- _all_tasks, _current_tasks)
+ _all_tasks, _current_tasks,
+ current_task)
 except ImportError:
 pass
 else:
+ _c_current_task = current_task
 _c_register_task = _register_task
 _c_unregister_task = _unregister_task
 _c_enter_task = _enter_task
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index 5168b8250ef0..e533d5273e9f 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -2804,6 +2804,7 @@ class CIntrospectionTests(test_utils.TestCase, BaseTaskIntrospectionTests):
 
 
 class BaseCurrentLoopTests:
+ current_task = None
 
 def setUp(self):
 super().setUp()
@@ -2814,33 +2815,39 @@ def new_task(self, coro):
 raise NotImplementedError
 
 def test_current_task_no_running_loop(self):
- self.assertIsNone(asyncio.current_task(loop=self.loop))
+ self.assertIsNone(self.current_task(loop=self.loop))
 
 def test_current_task_no_running_loop_implicit(self):
 with self.assertRaisesRegex(RuntimeError, 'no running event loop'):
- asyncio.current_task()
+ self.current_task()
 
 def test_current_task_with_implicit_loop(self):
 async def coro():
- self.assertIs(asyncio.current_task(loop=self.loop), task)
+ self.assertIs(self.current_task(loop=self.loop), task)
 
- self.assertIs(asyncio.current_task(None), task)
- self.assertIs(asyncio.current_task(), task)
+ self.assertIs(self.current_task(None), task)
+ self.assertIs(self.current_task(), task)
 
 task = self.new_task(coro())
 self.loop.run_until_complete(task)
- self.assertIsNone(asyncio.current_task(loop=self.loop))
+ self.assertIsNone(self.current_task(loop=self.loop))
 
 
 class PyCurrentLoopTests(BaseCurrentLoopTests, test_utils.TestCase):
+ current_task = staticmethod(tasks._py_current_task)
 
 def new_task(self, coro):
 return tasks._PyTask(coro, loop=self.loop)
 
 
- at unittest.skipUnless(hasattr(tasks, '_CTask'),
+ at unittest.skipUnless(hasattr(tasks, '_CTask') and
+ hasattr(tasks, '_c_current_task'),
 'requires the C _asyncio module')
 class CCurrentLoopTests(BaseCurrentLoopTests, test_utils.TestCase):
+ if hasattr(tasks, '_c_current_task'):
+ current_task = staticmethod(tasks._c_current_task)
+ else:
+ current_task = None
 
 def new_task(self, coro):
 return getattr(tasks, '_CTask')(coro, loop=self.loop)
diff --git a/Misc/NEWS.d/next/Library/2022-12-19-12-18-28.gh-issue-100344.lfCqpE.rst b/Misc/NEWS.d/next/Library/2022-12-19-12-18-28.gh-issue-100344.lfCqpE.rst
new file mode 100644
index 000000000000..d55f6888dbde
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2022-12-19-12-18-28.gh-issue-100344.lfCqpE.rst
@@ -0,0 +1,2 @@
+Provide C implementation for :func:`asyncio.current_task` for a 4x-6x
+speedup.
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
index 32be537c00a5..6fe4ca469475 100644
--- a/Modules/_asynciomodule.c
+++ b/Modules/_asynciomodule.c
@@ -3314,6 +3314,44 @@ _asyncio__leave_task_impl(PyObject *module, PyObject *loop, PyObject *task)
 }
 
 
+/*[clinic input]
+_asyncio.current_task
+
+ loop: object = None
+
+Return a currently executed task.
+
+[clinic start generated code]*/
+
+static PyObject *
+_asyncio_current_task_impl(PyObject *module, PyObject *loop)
+/*[clinic end generated code: output=fe15ac331a7f981a input=58910f61a5627112]*/
+{
+ PyObject *ret;
+ asyncio_state *state = get_asyncio_state(module);
+
+ if (loop == Py_None) {
+ loop = _asyncio_get_running_loop_impl(module);
+ if (loop == NULL) {
+ return NULL;
+ }
+ } else {
+ Py_INCREF(loop);
+ }
+
+ ret = PyDict_GetItemWithError(state->current_tasks, loop);
+ Py_DECREF(loop);
+ if (ret == NULL && PyErr_Occurred()) {
+ return NULL;
+ }
+ else if (ret == NULL) {
+ Py_RETURN_NONE;
+ }
+ Py_INCREF(ret);
+ return ret;
+}
+
+
 /*********************** Module **************************/
 
 
@@ -3494,6 +3532,7 @@ module_init(asyncio_state *state)
 PyDoc_STRVAR(module_doc, "Accelerator module for asyncio");
 
 static PyMethodDef asyncio_methods[] = {
+ _ASYNCIO_CURRENT_TASK_METHODDEF
 _ASYNCIO_GET_EVENT_LOOP_METHODDEF
 _ASYNCIO_GET_RUNNING_LOOP_METHODDEF
 _ASYNCIO__GET_RUNNING_LOOP_METHODDEF
diff --git a/Modules/clinic/_asynciomodule.c.h b/Modules/clinic/_asynciomodule.c.h
index f2fbb352c2c6..43c5d7717986 100644
--- a/Modules/clinic/_asynciomodule.c.h
+++ b/Modules/clinic/_asynciomodule.c.h
@@ -1242,4 +1242,64 @@ _asyncio__leave_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs,
 exit:
 return return_value;
 }
-/*[clinic end generated code: output=83580c190031241c input=a9049054013a1b77]*/
+
+PyDoc_STRVAR(_asyncio_current_task__doc__,
+"current_task($module, /, loop=None)\n"
+"--\n"
+"\n"
+"Return a currently executed task.");
+
+#define _ASYNCIO_CURRENT_TASK_METHODDEF \
+ {"current_task", _PyCFunction_CAST(_asyncio_current_task), METH_FASTCALL|METH_KEYWORDS, _asyncio_current_task__doc__},
+
+static PyObject *
+_asyncio_current_task_impl(PyObject *module, PyObject *loop);
+
+static PyObject *
+_asyncio_current_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
+{
+ PyObject *return_value = NULL;
+ #if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE)
+
+ #define NUM_KEYWORDS 1
+ static struct {
+ PyGC_Head _this_is_not_used;
+ PyObject_VAR_HEAD
+ PyObject *ob_item[NUM_KEYWORDS];
+ } _kwtuple = {
+ .ob_base = PyVarObject_HEAD_INIT(&PyTuple_Type, NUM_KEYWORDS)
+ .ob_item = { &_Py_ID(loop), },
+ };
+ #undef NUM_KEYWORDS
+ #define KWTUPLE (&_kwtuple.ob_base.ob_base)
+
+ #else // !Py_BUILD_CORE
+ # define KWTUPLE NULL
+ #endif // !Py_BUILD_CORE
+
+ static const char * const _keywords[] = {"loop", NULL};
+ static _PyArg_Parser _parser = {
+ .keywords = _keywords,
+ .fname = "current_task",
+ .kwtuple = KWTUPLE,
+ };
+ #undef KWTUPLE
+ PyObject *argsbuf[1];
+ Py_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - 0;
+ PyObject *loop = Py_None;
+
+ args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 0, 1, 0, argsbuf);
+ if (!args) {
+ goto exit;
+ }
+ if (!noptargs) {
+ goto skip_optional_pos;
+ }
+ loop = args[0];
+skip_optional_pos:
+ return_value = _asyncio_current_task_impl(module, loop);
+
+exit:
+ return return_value;
+}
+/*[clinic end generated code: output=00f494214f2fd008 input=a9049054013a1b77]*/
