Functional monkey-patched EpollSelector, good results

1 files changed, 11 insertions, 4 deletions

diff --git a/mini_scalene.py b/mini_scalene.py
index 569d7c0..f7a8ca9 100644
--- a/mini_scalene.py
+++ b/mini_scalene.py
@@ -1,3 +1,5 @@
+import selectors
+
 import sys
 import argparse
 import threading
@@ -13,6 +15,7 @@ from typing import (
 )
 from types import FrameType
 from collections import defaultdict
+import replacement_epoll_selector
 
 the_globals = {
     '__name__': '__main__',
@@ -64,7 +67,6 @@ class MiniScalene(object):
     profile_async = True
 
     def __init__(self):
-        import replacement_poll_selector
         signal.signal(signal.SIGPROF,
                       self.cpu_signal_handler)
         signal.setitimer(signal.ITIMER_PROF,
@@ -80,12 +82,16 @@ class MiniScalene(object):
     @staticmethod
     def start(profile_async):
         MiniScalene.profile_async = profile_async
-        atexit.register(MiniScalene.exit_handler)
 
     @staticmethod
-    def exit_handler():
-        '''Turn off profiling signals & pretty-print profiling information.'''
+    def stop():
+        '''Turn off profiling signals'''
         MiniScalene.disable_signals()
+        MiniScalene.exit_handler()
+
+    @staticmethod
+    def exit_handler():
+        '''Pretty-print profiling information.'''
         # If we've collected any samples, dump them.
         print("CPU usage (Python):")
         if MiniScalene.total_cpu_samples > 0:
@@ -250,6 +256,7 @@ def main():
             code = compile(fp.read(), args.script, "exec")
             MiniScalene().start(args.async_off)
             exec(code, the_globals)
+            MiniScalene().stop()
     except Exception:
         traceback.print_exc()