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

1 files changed, 12 insertions, 4 deletions

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