Support context in loop.create_task()

Author: fantix

uvloop/loop.pyx

@@ -50,6 +50,7 @@ include "errors.pyx"
 
 cdef:
     int PY39 = PY_VERSION_HEX >= 0x03090000
+    int PY311 = PY_VERSION_HEX >= 0x030b0000
     uint64_t MAX_SLEEP = 3600 * 24 * 365 * 100
 
 
@@ -1413,30 +1414,64 @@ cdef class Loop:
         """Create a Future object attached to the loop."""
         return self._new_future()
 
-    def create_task(self, coro, *, name=None):
-        """Schedule a coroutine object.
+    if PY311:
 
-        Return a task object.
+        def create_task(self, coro, *, name=None, context=None):
+            """Schedule a coroutine object.
 
-        If name is not None, task.set_name(name) will be called if the task
-        object has the set_name attribute, true for default Task in Python 3.8.
-        """
-        self._check_closed()
-        if self._task_factory is None:
-            task = aio_Task(coro, loop=self)
-        else:
-            task = self._task_factory(self, coro)
+            Return a task object.
 
-        # copied from asyncio.tasks._set_task_name (bpo-34270)
-        if name is not None:
-            try:
-                set_name = task.set_name
-            except AttributeError:
-                pass
+            If name is not None, task.set_name(name) will be called if the task
+            object has the set_name attribute, true for default Task in Python
+            3.8.
+
+            An optional keyword-only context argument allows specifying a
+            custom contextvars.Context for the coro to run in. The current
+            context copy is created when no context is provided.
+            """
+            self._check_closed()
+            if self._task_factory is None:
+                task = aio_Task(coro, loop=self, context=context)
+            else:
+                task = self._task_factory(self, coro, context=context)
+
+            # copied from asyncio.tasks._set_task_name (bpo-34270)
+            if name is not None:
+                try:
+                    set_name = task.set_name
+                except AttributeError:
+                    pass
+                else:
+                    set_name(name)
+
+            return task
+    else:
+
+        def create_task(self, coro, *, name=None):
+            """Schedule a coroutine object.
+
+            Return a task object.
+
+            If name is not None, task.set_name(name) will be called if the task
+            object has the set_name attribute, true for default Task in Python
+            3.8.
+            """
+            self._check_closed()
+            if self._task_factory is None:
+                task = aio_Task(coro, loop=self)
             else:
-                set_name(name)
+                task = self._task_factory(self, coro)
+
+            # copied from asyncio.tasks._set_task_name (bpo-34270)
+            if name is not None:
+                try:
+                    set_name = task.set_name
+                except AttributeError:
+                    pass
+                else:
+                    set_name(name)
 
-        return task
+            return task
 
     def set_task_factory(self, factory):
         """Set a task factory that will be used by loop.create_task().