Skip to content

gh-136157:Optimize asyncio.to_thread to avoid contextvars.copy_context() overhead for empty contexts #136158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
7 changes: 5 additions & 2 deletions Lib/asyncio/threads.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,5 +21,8 @@ async def to_thread(func, /, *args, **kwargs):
"""
loop = events.get_running_loop()
ctx = contextvars.copy_context()
func_call = functools.partial(ctx.run, func, *args, **kwargs)
return await loop.run_in_executor(None, func_call)
if len(ctx) == 0:
callback = functools.partial(func, *args, **kwargs)
else:
callback = functools.partial(ctx.run, func, *args, **kwargs)
return await loop.run_in_executor(None, callback)
38 changes: 37 additions & 1 deletion Lib/test/test_asyncio/test_threads.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,14 @@

import asyncio
import unittest
import functools

from contextvars import ContextVar
from unittest import mock


def tearDownModule():
asyncio._set_event_loop_policy(None)
asyncio.set_event_loop_policy(None)


class ToThreadTests(unittest.IsolatedAsyncioTestCase):
Expand Down Expand Up @@ -61,6 +62,41 @@ def get_ctx():

self.assertEqual(result, 'parrot')

@mock.patch('asyncio.base_events.BaseEventLoop.run_in_executor')
async def test_to_thread_optimization_path(self, run_in_executor):
# This test ensures that `to_thread` uses the correct execution path
# based on whether the context is empty or not.

# `to_thread` awaits the future returned by `run_in_executor`.
# We need to provide a completed future as a return value for the mock.
fut = asyncio.Future()
fut.set_result(None)
run_in_executor.return_value = fut

def myfunc():
pass

# Test with an empty context (optimized path)
await asyncio.to_thread(myfunc)
run_in_executor.assert_called_once()

callback = run_in_executor.call_args.args[1]
self.assertIsInstance(callback, functools.partial)
self.assertIs(callback.func, myfunc)
run_in_executor.reset_mock()

# Test with a non-empty context (standard path)
var = ContextVar('var')
var.set('value')

await asyncio.to_thread(myfunc)
run_in_executor.assert_called_once()

callback = run_in_executor.call_args.args[1]
self.assertIsInstance(callback, functools.partial)
self.assertIsNot(callback.func, myfunc) # Should be ctx.run
self.assertIs(callback.args[0], myfunc)


if __name__ == "__main__":
unittest.main()
40 changes: 31 additions & 9 deletions Lib/test/test_warnings/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
from test.support import import_helper
from test.support import os_helper
from test.support import warnings_helper
from test.support import script_helper
from _pyrepl.completing_reader import stripcolor as strip_ansi
from test.support import force_not_colorized
from test.support.script_helper import assert_python_ok, assert_python_failure

Expand Down Expand Up @@ -771,15 +773,35 @@ def test_improper_input(self):
self.assertRaises(UserWarning, self.module.warn, 'convert to error')

def test_import_from_module(self):
with self.module.catch_warnings():
self.module._setoption('ignore::Warning')
with self.assertRaises(self.module._OptionError):
self.module._setoption('ignore::TestWarning')
with self.assertRaises(self.module._OptionError):
self.module._setoption('ignore::test.test_warnings.bogus')
self.module._setoption('error::test.test_warnings.TestWarning')
with self.assertRaises(TestWarning):
self.module.warn('test warning', TestWarning)
with os_helper.temp_dir() as script_dir:
script = script_helper.make_script(script_dir,
'test_warnings_importer',
'import test.test_warnings.data.import_warning')
rc, out, err = assert_python_ok(script)
self.assertNotIn(b'UserWarning', err)

def test_syntax_warning_for_compiler(self):
# Test that SyntaxWarning from the compiler has a proper module name,
# not a guessed one like 'sys'. gh-135801
code = textwrap.dedent("""\
class A:
def func(self):
return self.var is 2
""")
# The name of the script is 'test_sw'
with os_helper.temp_dir() as script_dir:
script_name = script_helper.make_script(script_dir, 'test_sw', code)
# We want to check that the warning filter for 'test_sw' module works.
rc, out, err = assert_python_failure("-W", "error::SyntaxWarning:test_sw",
script_name)
self.assertEqual(rc, 1)
self.assertEqual(out, b'')
# Check that we got a SyntaxError.
err = err.decode()
err = strip_ansi(err)
self.assertIn("""SyntaxError: "is" with 'int' literal. Did you mean "=="?""", err)
# Check that the filename in the traceback is correct.
self.assertIn(os.path.basename(script_name), err)


class CWCmdLineTests(WCmdLineTests, unittest.TestCase):
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Fix warning deduplication regression for SyntaxWarnings with pseudo-filenames (e.g., <stdin>, <string>).
69 changes: 48 additions & 21 deletions Python/_warnings.c
Original file line number Diff line number Diff line change
Expand Up @@ -616,33 +616,60 @@ already_warned(PyInterpreterState *interp, PyObject *registry, PyObject *key,
static PyObject *
normalize_module(PyObject *filename)
{
PyObject *module;
int kind;
const void *data;
Py_ssize_t len;

len = PyUnicode_GetLength(filename);
if (len < 0)
PyObject *module_name = NULL;
PyObject *os_path = NULL;
PyObject *basename = NULL;
PyObject *splitext = NULL;
PyObject *root = NULL;

os_path = PyImport_ImportModule("os.path");
if (os_path == NULL) {
return NULL;
}

if (len == 0)
return PyUnicode_FromString("<unknown>");
basename = PyObject_CallMethod(os_path, "basename", "O", filename);
if (basename == NULL) {
goto cleanup;
}

kind = PyUnicode_KIND(filename);
data = PyUnicode_DATA(filename);
splitext = PyObject_CallMethod(os_path, "splitext", "O", basename);
if (splitext == NULL) {
goto cleanup;
}

/* if filename.endswith(".py"): */
if (len >= 3 &&
PyUnicode_READ(kind, data, len-3) == '.' &&
PyUnicode_READ(kind, data, len-2) == 'p' &&
PyUnicode_READ(kind, data, len-1) == 'y')
{
module = PyUnicode_Substring(filename, 0, len-3);
root = PyTuple_GetItem(splitext, 0);
if (root == NULL) {
goto cleanup;
}
else {
module = Py_NewRef(filename);

if (PyUnicode_CompareWithASCIIString(root, "__init__") == 0) {
PyObject *dirname = PyObject_CallMethod(os_path, "dirname", "O", filename);
if (dirname == NULL) {
goto cleanup;
}
module_name = PyObject_CallMethod(os_path, "basename", "O", dirname);
Py_DECREF(dirname);
} else {
module_name = Py_NewRef(root);
}

cleanup:
Py_XDECREF(os_path);
Py_XDECREF(basename);
Py_XDECREF(splitext);

if (module_name == NULL) {
// Fallback or error occurred
PyErr_Clear();
return PyUnicode_FromString("<unknown>");
}
return module;

if (PyUnicode_GetLength(module_name) == 0) {
Py_DECREF(module_name);
return PyUnicode_FromString("<unknown>");
}

return module_name;
}

static int
Expand Down
29 changes: 26 additions & 3 deletions Python/errors.c
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@

/* Error handling */

#include "Python.h"
Expand Down Expand Up @@ -1937,8 +1936,32 @@ int
_PyErr_EmitSyntaxWarning(PyObject *msg, PyObject *filename, int lineno, int col_offset,
int end_lineno, int end_col_offset)
{
if (_PyErr_WarnExplicitObjectWithContext(PyExc_SyntaxWarning, msg,
filename, lineno) < 0)
/* For pseudo-filenames (e.g., <string>, <stdin>), use the original approach
to maintain compatibility with existing behavior */
Py_ssize_t len = PyUnicode_GET_LENGTH(filename);
if (len > 1 &&
PyUnicode_READ_CHAR(filename, 0) == '<' &&
PyUnicode_READ_CHAR(filename, len - 1) == '>')
{
if (_PyErr_WarnExplicitObjectWithContext(PyExc_SyntaxWarning, msg,
filename, lineno) < 0)
{
if (PyErr_ExceptionMatches(PyExc_SyntaxWarning)) {
/* Replace the SyntaxWarning exception with a SyntaxError
to get a more accurate error report */
PyErr_Clear();
_PyErr_RaiseSyntaxError(msg, filename, lineno, col_offset,
end_lineno, end_col_offset);
}
return -1;
}
return 0;
}

/* For regular files, derive the module from the filename by passing NULL
as the module argument to PyErr_WarnExplicitObject */
if (PyErr_WarnExplicitObject(PyExc_SyntaxWarning, msg,
filename, lineno, NULL, NULL) < 0)
{
if (PyErr_ExceptionMatches(PyExc_SyntaxWarning)) {
/* Replace the SyntaxWarning exception with a SyntaxError
Expand Down
Loading