adding logger and propagate exit code option

Author: hnykda

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,8 @@ 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,
@@ -236,6 +242,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.