Add imperfect logic to include async generators in profiled frames

1 files changed, 10 insertions, 0 deletions

diff --git a/aergia/aergia.py b/aergia/aergia.py
index d14520a..6f43f2e 100755
--- a/aergia/aergia.py
+++ b/aergia/aergia.py
@@ -53,6 +53,7 @@ import sys
 import threading
 import time
 import traceback
+import inspect
 
 
 orig_thread_join = threading.Thread.join
@@ -247,6 +248,15 @@ class Aergia(object):
                 f = Aergia._get_deepest_traceable_frame(coro)
                 if f:
                     idle.append(f)
+
+        # handle async generators
+        # ideally, we would access these from _get_deepest_traceable_frame
+        # doing it this way causes us to assign generator time to the
+        # coroutine that called this generator in _get_deepest_traceable_frame
+        for ag in loop._asyncgens:
+            f = getattr(ag, 'ag_frame', None)
+            if f and Aergia._should_trace(f.f_code.co_filename):
+                idle.append(f)
         return idle
 
     @staticmethod