From 02ca06b3e835b34b577caf8931abe7c3786a7d87 Mon Sep 17 00:00:00 2001
From: bd <bdunahu@operationnull.com>
Date: Tue, 22 Jul 2025 21:01:35 -0600
Subject: Add imperfect logic to include async generators in profiled frames

---
 aergia/aergia.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

(limited to 'aergia')

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
-- 
cgit v1.2.3