task debug facilities -- set of all tasks, get and print stack

tulip/tasks.py

 """Support for tasks, coroutines and the scheduler."""
 
 __all__ = ['coroutine', 'Task',
            'FIRST_COMPLETED', 'FIRST_EXCEPTION', 'ALL_COMPLETED',
            'wait', 'wait_for', 'as_completed', 'sleep', 'async',
            'gather',
            ]
 
 import collections
 import concurrent.futures
 import functools
 import inspect
+import linecache
+import traceback
+import weakref
 
 from . import events
 from . import futures
 
 
 def coroutine(func):
     """Decorator to mark coroutines.
 
     Decorator wraps non generator functions and returns generator wrapper.
     If non generator function returns generator of Future it yield-from it.
@@ skipping 32 matching lines @@
 
     # An important invariant maintained while a Task not done:
     #
     # - Either _fut_waiter is None, and _step() is scheduled;
     # - or _fut_waiter is some Future, and _step() is *not* scheduled.
     #
     # The only transition from the latter to the former is through
     # _wakeup().  When _fut_waiter is not None, one of its callbacks
     # must be _wakeup().
 
+    # Weak set containing all tasks alive.
+    _all_tasks = weakref.WeakSet()
+
+    @classmethod
+    def all_tasks(cls):
+        """Return a set of all tasks in existence."""
+        return set(cls._all_tasks)
+
+    @classmethod
+    def all_pending_tasks(cls):
+        """Return a set of all tasks in existence that aren't done yet."""
+        return {t for t in cls._all_tasks if not t.done()}
+
+    @classmethod
+    def all_done_tasks(cls):
+        """Return a set of all tasks in existence that are done."""
+        return {t for t in cls._all_tasks if t.done()}
+
+    @classmethod
+    def all_failed_tasks(cls):
+        """Return a set of all tasks in existence that have failed.
+
+        This includes cancelled tasks and tasks with an exception set.
+        """
+        return {t for t in cls._all_tasks
+                  if t.done() and (t.cancelled() or t.exception())}
+
     def __init__(self, coro, *, loop=None):
         assert inspect.isgenerator(coro)  # Must be a coroutine *object*.
         super().__init__(loop=loop)
         self._coro = coro
         self._fut_waiter = None
         self._must_cancel = False
         self._loop.call_soon(self._step)
+        self.__class__._all_tasks.add(self)
 
     def __repr__(self):
         res = super().__repr__()
         if (self._must_cancel and
             self._state == futures._PENDING and
             '<PENDING' in res):
             res = res.replace('<PENDING', '<CANCELLING', 1)
         i = res.find('<')
         if i < 0:
             i = len(res)
         res = res[:i] + '(<{}>)'.format(self._coro.__name__) + res[i:]
         return res
 
+    def get_stack(self):
+        """Return the list of stack frames for this task's coroutine.
+
+        If the coroutine is active, this returns the stack where it is
+        suspended.  If the coroutine has completed successfully or was
+        cancelled, this returns an empty list.  If the coroutine was
+        terminated by an exception, this returns the list of traceback
+        frames.
+        """
+        frames = []
+        f = self._coro.gi_frame
+        if f is not None:
+            while f is not None:
+                frames.append(f)
+                f = f.f_back
+        elif self._exception is not None:
+            tb = self._exception.__traceback__
+            while tb is not None:
+                frames.append(tb.tb_frame)
+                tb = tb.tb_next
+            frames.reverse()
+        return frames
+
+    def print_stack(self, limit=None, file=None):
+        extracted_list = []
+        checked = set()
+        for f in self.get_stack():
+            lineno = f.f_lineno
+            co = f.f_code
+            filename = co.co_filename
+            name = co.co_name
+            if filename not in checked:
+                checked.add(filename)
+                linecache.checkcache(filename)
+            line = linecache.getline(filename, lineno, f.f_globals)
+            extracted_list.append((filename, lineno, name, line))
+        if not extracted_list:
+            print('No traceback for %r' % self, file=file)
+        elif self._exception is not None:
+            print('Traceback for %r (most recent call last):' % self,
+                  file=file)
+        else:
+            print('Stack for %r (most recent call last):' % self,
+                  file=file)
+        traceback.print_list(extracted_list, file=file)
+
     def cancel(self):
         if self.done():
             return False
         if self._fut_waiter is not None:
             if self._fut_waiter.cancel():
                 # Leave self._fut_waiter; it may be a Task that
                 # catches and ignores the cancellation so we may have
                 # to cancel it again later.
                 return True
         # It must be the case that self._step is already scheduled.
@@ skipping 368 matching lines @@
             outer.cancel()
         else:
             exc = inner.exception()
             if exc is not None:
                 outer.set_exception(exc)
             else:
                 outer.set_result(inner.result())
 
     inner.add_done_callback(_done_callback)
     return outer