Reformat repository to allow for new unit tests

2 files changed, 81 insertions, 52 deletions

diff --git a/aergia/__init__.py b/aergia/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/aergia/__init__.py
diff --git a/aergia b/aergia/aergia.py
index 3f629ef..b7b4f35 100755
--- a/aergia
+++ b/aergia/aergia.py
@@ -41,7 +41,7 @@ Commentary:
 Code:
 '''
 
-from collections import defaultdict
+from collections import defaultdict, namedtuple
 from typing import Optional
 import argparse
 import asyncio
@@ -78,6 +78,9 @@ def thread_join_replacement(
 
 threading.Thread.join = thread_join_replacement
 
+# a tuple used as a key in the sample-dict
+Sample = namedtuple('Sample', ['file', 'line', 'func'])
+
 
 class Aergia(object):
 
@@ -88,9 +91,11 @@ class Aergia(object):
     # number of times samples have been collected
     total_samples = 0
     # the (ideal) interval between samples
+
     signal_interval = 0.0
 
-    def __init__(self, signal_interval):
+    @staticmethod
+    def __init__(signal_interval):
         Aergia.signal_interval = signal_interval
 
     @staticmethod
@@ -111,8 +116,18 @@ class Aergia(object):
 
     @staticmethod
     def stop():
-        Aergia.disable_signals()
-        Aergia.print_samples()
+        '''Stops the profiler.'''
+        signal.setitimer(signal.ITIMER_REAL, 0)
+
+    @staticmethod
+    def clear():
+        Aergia.total_samples = 0
+        Aergia.samples = defaultdict(lambda: 0)
+
+    @staticmethod
+    def get_samples():
+        '''Returns the profiling results.'''
+        return Aergia.samples
 
     @staticmethod
     def print_samples():
@@ -130,20 +145,17 @@ class Aergia(object):
         '''Pretty-print a single sample.'''
         sig_intv = Aergia.signal_interval
         value = Aergia.samples[key]
-        print(f"{key} :\t\t{value * 100 / Aergia.total_samples:.3f}%"
+        print(f"{Aergia.tuple_to_string(key)} :"
+              f"\t\t{value * 100 / Aergia.total_samples:.3f}%"
               f"\t({value:.3f} ->"
               f" {value*sig_intv:.6f} seconds)")
 
     @staticmethod
-    def disable_signals():
-        signal.setitimer(signal.ITIMER_REAL, 0)
-
-    @staticmethod
     def idle_signal_handler(sig, frame):
         '''Obtains and records which lines are currently being waited on.'''
         keys = Aergia.compute_frames_to_record()
         for key in keys:
-            Aergia.samples[Aergia.frame_to_string(key)] += 1
+            Aergia.samples[Aergia.frame_to_tuple(key)] += 1
             Aergia.total_samples += 1
 
     @staticmethod
@@ -159,10 +171,7 @@ class Aergia(object):
         what is running for us.'''
         loops = Aergia.get_event_loops()
         frames = Aergia.get_frames_from_loops(loops)
-        return [
-            f for f in frames
-            if f is not None and Aergia.should_trace(f.filename)
-        ]
+        return frames
 
     @staticmethod
     def get_event_loops():
@@ -206,28 +215,18 @@ class Aergia(object):
         ]
 
     @staticmethod
-    def frame_to_string(frame):
-        '''Pretty-prints a frame as a function/file name and a line number.
-        Additionally used as a key for tallying lines.'''
-        func_name = frame.name
-        line_no = frame.lineno
-        filename = frame.filename
-        return filename + ':' + str(line_no) + '\t' + func_name
-
-    @staticmethod
-    def should_trace(filename):
-        '''Returns FALSE if filename is uninteresting to the user.'''
-        # return True
-        # FIXME Assume GuixSD. Makes filtering easy
-        if '/gnu/store' in filename:
-            return False
-        if 'site-packages' in filename:
-            return False
-        if filename[0] == '<':
-            return False
-        if 'aergia' in filename:
-            return False
-        return True
+    def frame_to_tuple(frame):
+        '''Given a frame, constructs a sample key for tallying lines.'''
+        co = frame.f_code
+        func_name = co.co_name
+        line_no = frame.f_lineno
+        filename = co.co_filename
+        return Sample(filename, line_no, func_name)
+
+    def tuple_to_string(sample):
+        '''Given a namedtuple corresponding to a sample key,
+        pretty-prints a frame as a function/file name and a line number.'''
+        return sample.file + ':' + str(sample.line) + '\t' + sample.func
 
     @staticmethod
     def sort_samples(sample_dict):
@@ -241,13 +240,50 @@ 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 = []
-        for th in loop._scheduled:
-            st = th._source_traceback
-            if st:
-                idle += st
+        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)
+                if f:
+                    idle.append(f)
         return idle
 
     @staticmethod
+    def get_deepest_traceable_frame(coro):
+        if not coro:
+            return None
+        curr = coro
+        lframe = None
+        while True:
+            frame = getattr(curr, 'cr_frame', None)
+            if not frame or not Aergia.should_trace(frame.f_code.co_filename):
+                return lframe
+
+            lframe = frame
+            awaited = getattr(curr, 'cr_await', None)
+            if not awaited or not hasattr(awaited, 'cr_frame'):
+                return lframe
+            curr = awaited
+
+    @staticmethod
+    def should_trace(filename):
+        '''Returns FALSE if filename is uninteresting to the user.'''
+        # print(filename)
+        # FIXME Assume GuixSD. Makes filtering easy
+        if '/gnu/store' in filename:
+            return False
+        if 'site-packages' in filename:
+            return False
+        if filename[0] == '<':
+            return False
+        # if 'aergia' in filename:
+        #     return False
+        return True
+
+    @staticmethod
     def gettime():
         '''returns the wallclock time'''
         return time.process_time()
@@ -266,8 +302,9 @@ the_globals = {
 }
 
 
-def parse_arguments():
-    '''Parse CLI args'''
+if __name__ == "__main__":
+    # parses CLI arguments and facilitates profiler runtime.
+    # foo
     parser = argparse.ArgumentParser(
         usage='%(prog)s [args] script [args]'
     )
@@ -281,12 +318,7 @@ def parse_arguments():
     parser.add_argument('script', help='A python script to run.')
     parser.add_argument('s_args', nargs=argparse.REMAINDER,
                         help='python script args')
-
-    return parser.parse_args()
-
-
-def main():
-    args = parse_arguments()
+    args = parser.parse_args()
 
     sys.argv = [args.script] + args.s_args
     try:
@@ -294,10 +326,7 @@ def main():
             code = compile(fp.read(), args.script, "exec")
             Aergia(args.interval).start()
             exec(code, the_globals)
+            Aergia.print_samples()
             Aergia.stop()
     except Exception:
         traceback.print_exc()
-
-
-if __name__ == "__main__":
-    main()