From 8ad9907ba36d8a93fae61be97a4fcd7dfac4ad91 Mon Sep 17 00:00:00 2001
From: bd <bdunahu@operationnull.com>
Date: Mon, 16 Jun 2025 00:23:37 -0400
Subject: Add LICENSE, docstring, and small cleanups to aergia.py

---
 aergia.py | 142 +++++++++++++++++++++++++++++++++++++++++++++++---------------
 1 file changed, 109 insertions(+), 33 deletions(-)

(limited to 'aergia.py')

diff --git a/aergia.py b/aergia.py
index 534880f..64567a2 100644
--- a/aergia.py
+++ b/aergia.py
@@ -1,47 +1,85 @@
-import sys
+'''
+         _/_/                                  _/
+      _/    _/    _/_/    _/  _/_/    _/_/_/        _/_/_/
+     _/_/_/_/  _/_/_/_/  _/_/      _/    _/  _/  _/    _/
+    _/    _/  _/        _/        _/    _/  _/  _/    _/
+   _/    _/    _/_/_/  _/          _/_/_/  _/    _/_/_/
+                                      _/
+                                 _/_/
+Copyright:
+
+   This program is free software: you can redistribute it
+   and/or modify it under the terms of the GNU General
+   Public License as published by the Free Software
+   Foundation, either version 3 of the License, or (at your
+   option) any later version.
+
+   This program is distributed in the hope that it will be
+   useful, but WITHOUT ANY WARRANTY; without even the implied
+   warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
+   PURPOSE. See the GNU General Public License for more details.
+
+   You should have received a copy of the GNU General Public
+   License along with this program. If not, see
+   <https://www.gnu.org/licenses/>.
+
+
+Commentary:
+
+   Aergia is a sampling based profiler based off of SCALENE
+   (https://github.com/plasma-umass/scalene).
+
+   It is not particularly informative, but unlike SCALENE
+   or other sampling-based profilers I could find, reports
+   the wall-time each asyncio await call spends idling.
+
+   The goal behind Aergia is to eventually have these
+   features, or similar, merged into SCALENE.
+
+
+Code:
+'''
+
+from collections import defaultdict
+from types import FrameType
+from typing import cast, List, Tuple
 import argparse
-import threading
-import traceback
-import signal
 import asyncio
+import selectors
+import signal
+import sys
+import threading
 import time
-from typing import cast
-from types import FrameType
-from collections import defaultdict
-import replacement_epoll_selector
-
-the_globals = {
-    '__name__': '__main__',
-    '__doc__': None,
-    '__package__': None,
-    '__loader__': globals()['__loader__'],
-    '__spec__': None,
-    '__annotations__': {},
-    '__builtins__': globals()['__builtins__'],
-    '__file__': None,
-    '__cached__': None,
-}
+import traceback
 
 
-def parse_arguments():
-    '''Parse CLI args'''
-    parser = argparse.ArgumentParser()
+class ReplacementEpollSelector(selectors.EpollSelector):
+    '''
+    Provides a replacement for selectors.PollSelector that
+    periodically wakes up to accept signals.
+    '''
 
-    parser.add_argument('-a', '--async_off',
-                        action='store_false',
-                        help='Turn off experimental async profiling.',
-                        default=True)
+    def select(
+        self, timeout=None
+    ) -> List[Tuple[selectors.SelectorKey, int]]:
+        start_time = time.perf_counter()
+        if not timeout or timeout < 0:
+            interval = sys.getswitchinterval()
+        else:
+            interval = min(timeout, sys.getswitchinterval())
+        while True:
+            selected = super().select(interval)
+            if selected or timeout == 0 or not timeout:
+                return selected
+            end_time = time.perf_counter()
+            if end_time - start_time >= timeout:
+                return []
 
-    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()
+selectors.DefaultSelector = ReplacementEpollSelector
 
 
 class Aergia(object):
-    '''A stripped-down version of SCALENE which tallies active lines during
-    execution.'''
 
     # a key-value pair where keys represent frame metadata (see
     # Aergia.frame_to_string) and values represent number of times
@@ -221,6 +259,42 @@ class Aergia(object):
         return list(s)
 
 
+the_globals = {
+    '__name__': '__main__',
+    '__doc__': None,
+    '__package__': None,
+    '__loader__': globals()['__loader__'],
+    '__spec__': None,
+    '__annotations__': {},
+    '__builtins__': globals()['__builtins__'],
+    '__file__': None,
+    '__cached__': None,
+}
+
+
+def parse_arguments():
+    '''Parse CLI args'''
+    parser = argparse.ArgumentParser(
+        usage='%(prog)s [args] script [args]'
+    )
+
+    parser.add_argument('-a', '--async_off',
+                        action='store_false',
+                        help='Turn off experimental async profiling.',
+                        default=True)
+
+    parser.add_argument('-i', '--interval',
+                        help='The minimum amount of time inbetween \
+                        samples in seconds.',
+                        metavar='',
+                        default=0.01)
+    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()
 
@@ -237,3 +311,5 @@ def main():
 
 if __name__ == "__main__":
     main()
+
+# aergia.py ends here
-- 
cgit v1.2.3