Profile the current task as well

2 files changed, 9 insertions, 7 deletions

diff --git a/aergia/aergia.py b/aergia/aergia.py
index 6f43f2e..28e8d7a 100755
--- a/aergia/aergia.py
+++ b/aergia/aergia.py
@@ -53,7 +53,6 @@ import sys
 import threading
 import time
 import traceback
-import inspect
 
 
 orig_thread_join = threading.Thread.join
@@ -110,7 +109,6 @@ class Aergia(object):
         in `Aergia._get_idle_task_frames'. The TimerHandler object otherwise
         does not keep track of a _source_traceback.
         '''
-        os.environ["PYTHONASYNCIODEBUG"] = "1"
         signal.signal(signal.SIGALRM,
                       Aergia._idle_signal_handler)
         signal.setitimer(signal.ITIMER_REAL,
@@ -239,10 +237,7 @@ class Aergia(object):
         '''Given an asyncio event loop, returns the list of idle task frames.
         A task is considered 'idle' if it is not currently executing.'''
         idle = []
-        current = asyncio.current_task(loop)
         for task in asyncio.all_tasks(loop):
-            if task == current:
-                continue
             coro = task.get_coro()
             if coro:
                 f = Aergia._get_deepest_traceable_frame(coro)
@@ -316,12 +311,19 @@ if __name__ == "__main__":
                         metavar='',
                         type=float,
                         default=0.01)
+    parser.add_argument('-d', '--debug',
+                        help='Turn on debug information for the event loop.',
+                        metavar='',
+                        type=bool,
+                        default=False)
     parser.add_argument('script', help='A python script to run.')
     parser.add_argument('s_args', nargs=argparse.REMAINDER,
                         help='python script args')
     args = parser.parse_args()
 
     sys.argv = [args.script] + args.s_args
+    if args.debug:
+        os.environ["PYTHONASYNCIODEBUG"] = "1"
     try:
         with open(args.script, 'r', encoding='utf-8') as fp:
             code = compile(fp.read(), args.script, "exec")
diff --git a/t/test_functionality.py b/t/test_functionality.py
index eb9873a..32e89f8 100644
--- a/t/test_functionality.py
+++ b/t/test_functionality.py
@@ -240,8 +240,8 @@ class BasicUsage(utils.AergiaUnitTestCase):
         self.assert_reasonable_delay('a', 0, aergia_samples)
 
         self.assert_similar_delay('c', yappi_samples, aergia_samples)
-        # Aergia does not assign time to the current task.
-        # These should fail.
+        # Aergia only assigns time to the current line where the task is
+        # suspended, not the whole coroutine stack.
         # self.assert_similar_delay('b', yappi_samples, aergia_samples)
         # self.assert_similar_delay('a', yappi_samples, aergia_samples)