bpo-34270: Make it possible to name asyncio tasks (GH-8547)

Co-authored-by: Antti Haapala <antti.haapala@anttipatterns.com>

4 files changed, 41 insertions, 8 deletions

diff --git a/Lib/asyncio/base_events.py b/Lib/asyncio/base_events.py
index 78fe2a7..ee13d1a 100644
--- a/Lib/asyncio/base_events.py
+++ b/Lib/asyncio/base_events.py
@@ -384,18 +384,20 @@ class BaseEventLoop(events.AbstractEventLoop):
         """Create a Future object attached to the loop."""
         return futures.Future(loop=self)
 
-    def create_task(self, coro):
+    def create_task(self, coro, *, name=None):
         """Schedule a coroutine object.
 
         Return a task object.
         """
         self._check_closed()
         if self._task_factory is None:
-            task = tasks.Task(coro, loop=self)
+            task = tasks.Task(coro, loop=self, name=name)
             if task._source_traceback:
                 del task._source_traceback[-1]
         else:
             task = self._task_factory(self, coro)
+            tasks._set_task_name(task, name)
+
         return task
 
     def set_task_factory(self, factory):
diff --git a/Lib/asyncio/base_tasks.py b/Lib/asyncio/base_tasks.py
index 3ce51f6..e2da462 100644
--- a/Lib/asyncio/base_tasks.py
+++ b/Lib/asyncio/base_tasks.py
@@ -12,11 +12,13 @@ def _task_repr_info(task):
         # replace status
         info[0] = 'cancelling'
 
+    info.insert(1, 'name=%r' % task.get_name())
+
     coro = coroutines._format_coroutine(task._coro)
-    info.insert(1, f'coro=<{coro}>')
+    info.insert(2, f'coro=<{coro}>')
 
     if task._fut_waiter is not None:
-        info.insert(2, f'wait_for={task._fut_waiter!r}')
+        info.insert(3, f'wait_for={task._fut_waiter!r}')
     return info
 
 
diff --git a/Lib/asyncio/events.py b/Lib/asyncio/events.py
index e4e6322..58a60a0 100644
--- a/Lib/asyncio/events.py
+++ b/Lib/asyncio/events.py
@@ -277,7 +277,7 @@ class AbstractEventLoop:
 
     # Method scheduling a coroutine object: create a task.
 
-    def create_task(self, coro):
+    def create_task(self, coro, *, name=None):
         raise NotImplementedError
 
     # Methods for interacting with threads.
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 72792a2..03d71d3 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -13,6 +13,7 @@ import concurrent.futures
 import contextvars
 import functools
 import inspect
+import itertools
 import types
 import warnings
 import weakref
@@ -23,6 +24,11 @@ from . import events
 from . import futures
 from .coroutines import coroutine
 
+# Helper to generate new task names
+# This uses itertools.count() instead of a "+= 1" operation because the latter
+# is not thread safe. See bpo-11866 for a longer explanation.
+_task_name_counter = itertools.count(1).__next__
+
 
 def current_task(loop=None):
     """Return a currently executed task."""
@@ -48,6 +54,16 @@ def _all_tasks_compat(loop=None):
     return {t for t in _all_tasks if futures._get_loop(t) is loop}
 
 
+def _set_task_name(task, name):
+    if name is not None:
+        try:
+            set_name = task.set_name
+        except AttributeError:
+            pass
+        else:
+            set_name(name)
+
+
 class Task(futures._PyFuture):  # Inherit Python Task implementation
                                 # from a Python Future implementation.
 
@@ -94,7 +110,7 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
                       stacklevel=2)
         return _all_tasks_compat(loop)
 
-    def __init__(self, coro, *, loop=None):
+    def __init__(self, coro, *, loop=None, name=None):
         super().__init__(loop=loop)
         if self._source_traceback:
             del self._source_traceback[-1]
@@ -104,6 +120,11 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
             self._log_destroy_pending = False
             raise TypeError(f"a coroutine was expected, got {coro!r}")
 
+        if name is None:
+            self._name = f'Task-{_task_name_counter()}'
+        else:
+            self._name = str(name)
+
         self._must_cancel = False
         self._fut_waiter = None
         self._coro = coro
@@ -126,6 +147,12 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
     def _repr_info(self):
         return base_tasks._task_repr_info(self)
 
+    def get_name(self):
+        return self._name
+
+    def set_name(self, value):
+        self._name = str(value)
+
     def set_result(self, result):
         raise RuntimeError('Task does not support set_result operation')
 
@@ -312,13 +339,15 @@ else:
     Task = _CTask = _asyncio.Task
 
 
-def create_task(coro):
+def create_task(coro, *, name=None):
     """Schedule the execution of a coroutine object in a spawn task.
 
     Return a Task object.
     """
     loop = events.get_running_loop()
-    return loop.create_task(coro)
+    task = loop.create_task(coro)
+    _set_task_name(task, name)
+    return task
 
 
 # wait() and as_completed() similar to those in PEP 3148.