replace profiling script with context manager

amodem/profiling.py

@@ -0,0 +1,14 @@
+import cProfile
+import contextlib
+
+
+@contextlib.contextmanager
+def save(filename=None):
+    if filename:
+        pr = cProfile.Profile()
+        pr.enable()
+        yield
+        pr.disable()
+        pr.dump_stats(filename)
+    else:
+        yield

amodem/recv.py

@@ -18,6 +18,7 @@ from . import train
 from . import common
 from . import config
 from . import ecc
+from . import profiling
 
 modem = sigproc.MODEM(config)
 
@@ -290,9 +291,11 @@ if __name__ == '__main__':
                    default=sys.stdin)
     p.add_argument('-o', '--output', type=argparse.FileType('wb'),
                    default=sys.stdout)
+    p.add_argument('-p', '--profile', type=str)
     args = p.parse_args()
     try:
-        main(args)
+        with profiling.save(filename=args.profile):
+            main(args)
     except Exception as e:
         log.exception(e)
     finally:

scripts/profile.sh

@@ -1,5 +0,0 @@
-#!/bin/bash
-python -m cProfile -o result.prof $*
-
-echo "sort time
-stats 20" | python -m pstats result.prof 1>&2