bpo-46994: Accept explicit contextvars.Context in asyncio create_task() API (GH-31837)

4 files changed, 26 insertions, 10 deletions

diff --git a/Lib/asyncio/base_events.py b/Lib/asyncio/base_events.py
index 51c4e66..5eea165 100644
--- a/Lib/asyncio/base_events.py
+++ b/Lib/asyncio/base_events.py
@@ -426,18 +426,23 @@ class BaseEventLoop(events.AbstractEventLoop):
         """Create a Future object attached to the loop."""
         return futures.Future(loop=self)
 
-    def create_task(self, coro, *, name=None):
+    def create_task(self, coro, *, name=None, context=None):
         """Schedule a coroutine object.
 
         Return a task object.
         """
         self._check_closed()
         if self._task_factory is None:
-            task = tasks.Task(coro, loop=self, name=name)
+            task = tasks.Task(coro, loop=self, name=name, context=context)
             if task._source_traceback:
                 del task._source_traceback[-1]
         else:
-            task = self._task_factory(self, coro)
+            if context is None:
+                # Use legacy API if context is not needed
+                task = self._task_factory(self, coro)
+            else:
+                task = self._task_factory(self, coro, context=context)
+
             tasks._set_task_name(task, name)
 
         return task
diff --git a/Lib/asyncio/events.py b/Lib/asyncio/events.py
index e682a19..0d26ea5 100644
--- a/Lib/asyncio/events.py
+++ b/Lib/asyncio/events.py
@@ -274,7 +274,7 @@ class AbstractEventLoop:
 
     # Method scheduling a coroutine object: create a task.
 
-    def create_task(self, coro, *, name=None):
+    def create_task(self, coro, *, name=None, context=None):
         raise NotImplementedError
 
     # Methods for interacting with threads.
diff --git a/Lib/asyncio/taskgroups.py b/Lib/asyncio/taskgroups.py
index c3ce94a..6af21f3 100644
--- a/Lib/asyncio/taskgroups.py
+++ b/Lib/asyncio/taskgroups.py
@@ -138,12 +138,15 @@ class TaskGroup:
             me = BaseExceptionGroup('unhandled errors in a TaskGroup', errors)
             raise me from None
 
-    def create_task(self, coro, *, name=None):
+    def create_task(self, coro, *, name=None, context=None):
         if not self._entered:
             raise RuntimeError(f"TaskGroup {self!r} has not been entered")
         if self._exiting and self._unfinished_tasks == 0:
             raise RuntimeError(f"TaskGroup {self!r} is finished")
-        task = self._loop.create_task(coro)
+        if context is None:
+            task = self._loop.create_task(coro)
+        else:
+            task = self._loop.create_task(coro, context=context)
         tasks._set_task_name(task, name)
         task.add_done_callback(self._on_task_done)
         self._unfinished_tasks += 1
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index e604298..b4f1eed 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -93,7 +93,7 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
     # status is still pending
     _log_destroy_pending = True
 
-    def __init__(self, coro, *, loop=None, name=None):
+    def __init__(self, coro, *, loop=None, name=None, context=None):
         super().__init__(loop=loop)
         if self._source_traceback:
             del self._source_traceback[-1]
@@ -112,7 +112,10 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
         self._must_cancel = False
         self._fut_waiter = None
         self._coro = coro
-        self._context = contextvars.copy_context()
+        if context is None:
+            self._context = contextvars.copy_context()
+        else:
+            self._context = context
 
         self._loop.call_soon(self.__step, context=self._context)
         _register_task(self)
@@ -360,13 +363,18 @@ else:
     Task = _CTask = _asyncio.Task
 
 
-def create_task(coro, *, name=None):
+def create_task(coro, *, name=None, context=None):
     """Schedule the execution of a coroutine object in a spawn task.
 
     Return a Task object.
     """
     loop = events.get_running_loop()
-    task = loop.create_task(coro)
+    if context is None:
+        # Use legacy API if context is not needed
+        task = loop.create_task(coro)
+    else:
+        task = loop.create_task(coro, context=context)
+
     _set_task_name(task, name)
     return task