1 files changed, 73 insertions, 21 deletions

diff --git a/aergia/aergia.py b/aergia/aergia.py
index 9b72d2f..48efc68 100755
--- a/aergia/aergia.py
+++ b/aergia/aergia.py
@@ -266,35 +266,91 @@ class Aergia(object):
     @staticmethod
     def _get_idle_task_frames(loop):
         '''Given an asyncio event loop, returns the list of idle task frames.
-        A task is considered 'idle' if it is not currently executing.'''
+        We only care about idle task frames, as running tasks are already
+        included elsewhere.
+
+        A task is considered 'idle' if it is pending and not the current
+        task.'''
         idle = []
+        coros = Aergia._get_traceable_coros(loop)
+        for coro in coros:
+            frame = Aergia._get_deepest_traceable_frame(coro)
+            # if it's a future,
+            if frame:
+                idle.append(frame)
+
+        # 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
+    def _get_traceable_coros(loop):
+        '''Given event loop LOOP, returns a list of coroutines corresponding
+        to each traceable task.
+
+        A task is traceable if it is not the current task, if it has actually
+        started executing, and if a child task did not originate from it.'''
         current = asyncio.current_task(loop)
+
+        coros = dict()
+        p_tasks = set()
         for task in asyncio.all_tasks(loop):
             # the task is not idle
             if task == current:
                 continue
+
+            coro = task.get_coro()
+
             # the task hasn't even run yet
             # assumes that all started tasks are sitting at an await
             # statement.
             # if this isn't the case, the associated coroutine will
             # be 'waiting' on the coroutine declaration. No! Bad!
-            if getattr(task.get_coro(), 'cr_await', None) is None:
+            if getattr(coro, 'cr_frame', None) is None or \
+               getattr(coro, 'cr_await', None) is None:
                 continue
-            coro = task.get_coro()
-            if coro:
-                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
+            frame = getattr(coro, 'cr_frame', None)
+
+            # for methods like gather, tasks create new tasks. keep
+            # track of these tasks so we don't count them twice.
+            parent = Aergia._get_parent_task(task)
+            if parent is not None:
+                p_tasks.add(parent)
+
+            coros[coro] = (frame.f_code.co_filename,
+                           frame.f_lineno,
+                           frame.f_code.co_name)
+
+        return [
+            coro for coro, frame_summary in coros.items()
+            if frame_summary not in p_tasks
+        ]
+
+    @staticmethod
+    def _get_parent_task(task):
+        '''Given TASK, returns the first frame in the source traceback
+        which is tracable. This only works if the event loop has debug
+        mode turned on.
+
+        If this is not the case, returns None.'''
+        tb = task._source_traceback
+
+        # true when debug mode isn't turned on
+        if tb is None:
+            return None
+
+        for frame in tb[::-1]:
+            if Aergia._should_trace(frame.filename):
+                return (frame.filename, frame.lineno, frame.name)
+
+        return None
 
     @staticmethod
     def _get_deepest_traceable_frame(coro):
@@ -373,9 +429,6 @@ if __name__ == "__main__":
                         metavar='',
                         type=float,
                         default=0.01)
-    parser.add_argument('-d', '--debug',
-                        help='Turn on debug information for the event loop.',
-                        action='store_true')
     parser.add_argument('-a', '--async-only',
                         help='Do not profile currently running tasks.',
                         action='store_true')
@@ -385,8 +438,7 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     sys.argv = [args.script] + args.s_args
-    if args.debug:
-        os.environ["PYTHONASYNCIODEBUG"] = "1"
+    os.environ["PYTHONASYNCIODEBUG"] = "1"
     try:
         with open(args.script, 'r', encoding='utf-8') as fp:
             code = compile(fp.read(), args.script, "exec")