Merge pull request #220 from hnykda/master

Optionally propagate the exit code

Author: astrojuanlu

Makefile

@@ -16,6 +16,7 @@ test:
 	$(PYTHON) test/test_memory_usage.py
 	$(PYTHON) test/test_precision_import.py
 	$(PYTHON) test/test_exception.py
+	$(PYTHON) test/test_exit_code.py
 	$(PYTHON) test/test_mprof.py
 
 develop:

mprof.py

@@ -6,6 +6,7 @@
 import copy
 import time
 import math
+import logging
 
 from collections import defaultdict
 from argparse import ArgumentParser, ArgumentError, REMAINDER, RawTextHelpFormatter
@@ -26,6 +27,9 @@
 For example, mprof plot --help will list all plotting options.
 """
 
+logger = logging.getLogger(__name__)
+logging.basicConfig()
+
 
 def print_usage():
     print("Usage: %s <command> <options> <arguments>"
@@ -185,6 +189,7 @@ def run_action():
                         help="""Monitors forked processes as well (sum up all process memory)""")
     parser.add_argument("--multiprocess", "-M", dest="multiprocess", action="store_true",
                         help="""Monitors forked processes creating individual plots for each child (disables --python features)""")
+    parser.add_argument("--exit-code", "-E", dest="exit_code", action="store_true", help="""Propagate the exit code""")
     parser.add_argument("--output", "-o", dest="filename",
                         default="mprofile_%s.dat" % time.strftime("%Y%m%d%H%M%S", time.localtime()),
                         help="""File to store results in, defaults to 'mprofile_<YYYYMMDDhhmmss>.dat' in the current directory,
@@ -197,7 +202,7 @@ def run_action():
                              'Option 4: (--python flag present) "<PYTHON_MODULE> <ARG1> <ARG2>..." - profile python module\n'
                         )
     args = parser.parse_args()
-
+    
     if len(args.program) == 0:
         print("A program to run must be provided. Use -h for help")
         sys.exit(1)
@@ -236,6 +241,11 @@ def run_action():
                         include_children=args.include_children,
                         multiprocess=args.multiprocess, stream=f)
 
+    if args.exit_code:
+        if p.returncode != 0:
+            logger.error('Program resulted with a non-zero exit code: %s', p.returncode)
+        sys.exit(p.returncode)
+
 
 def add_brackets(xloc, yloc, xshift=0, color="r", label=None, options=None):
     """Add two brackets on the memory line plot.

test/test_exit_code.py

@@ -0,0 +1,42 @@
+import unittest
+import sys
+import tempfile
+
+
+class TestExitCode(unittest.TestCase):
+
+    def setUp(self):
+        # to be able to import mprof
+        sys.path.append('.')
+        from mprof import run_action
+        self.run_action = run_action
+
+    def test_exit_code_success(self):
+        s = "1+1"
+        tmpfile = tempfile.NamedTemporaryFile('w', suffix='.py')
+        with tmpfile as ofile:
+            ofile.write(s)
+            ofile.flush()
+            sys.argv = ['<ignored>', '--exit-code', tmpfile.name]
+            self.assertRaisesRegexp(SystemExit, '0', self.run_action)
+
+    def test_exit_code_fail(self):
+        s = "raise RuntimeError('I am not working nicely')"
+        tmpfile = tempfile.NamedTemporaryFile('w', suffix='.py')
+        with tmpfile as ofile:
+            ofile.write(s)
+            ofile.flush()
+            sys.argv = ['<ignored>', '--exit-code', tmpfile.name]
+            self.assertRaisesRegexp(SystemExit, '1', self.run_action)
+
+    def test_no_exit_code_success(self):
+        s = "raise RuntimeError('I am not working nicely')"
+        tmpfile = tempfile.NamedTemporaryFile('w', suffix='.py')
+        with tmpfile as ofile:
+            ofile.write(s)
+            ofile.flush()
+            sys.argv = ['<ignored>', tmpfile.name]
+            self.run_action()
+
+if __name__ == '__main__':
+    unittest.main()