Plot improvements

mprof

@@ -157,6 +157,12 @@ def clean_action():
             os.remove(filename)
 
 
+def get_cmd_line(args):
+    """Given a set or arguments, compute command-line."""
+    blanks = set(' \t')
+    args = [s if blanks.isdisjoint(s) else "'" + s + "'" for s in args]
+    return ' '.join(args)
+
 
 def run_action():
     import time, subprocess
@@ -200,15 +206,18 @@ def run_action():
         print("running as a Python program...")
         if not args[0].startswith("python"):
             args.insert(0, "python")
+        cmd_line = get_cmd_line(args)
         args[1:1] = ("-m", "memory_profiler", "--timestamp",
                      "-o", mprofile_output)
         p = subprocess.Popen(args)
     else:
+        cmd_line = get_cmd_line(args)
         p = subprocess.Popen(args)
 
-    mu = mp.memory_usage(proc=p, interval=options.interval, timestamps=True,
-                         include_children=options.include_children)
     with open(mprofile_output, "a") as f:
+        f.write("CMDLINE {0}\n".format(cmd_line))
+        mu = mp.memory_usage(proc=p, interval=options.interval, timestamps=True,
+                         include_children=options.include_children)
         for m, t in mu:
             f.write("MEM {0:.6f} {1:.4f}".format(m, t) + "\n")
 
@@ -259,31 +268,50 @@ def add_brackets(xloc, yloc, xshift=0, color="r", label=None):
 
 
 def read_mprofile_file(filename):
-    # TODO: would be nice to do without numpy
+    """Read an mprofile file and return its content.
 
-    import numpy as np
-    ret = {}
-    mdata = []
+    Returns
+    =======
+    content: dict
+        Keys:
+
+        - "mem_usage": (list) memory usage values, in MiB
+        - "timestamp": (list) time instant for each memory usage value, in
+            second
+        - "func_timestamp": (dict) for each function, timestamps and memory
+            usage upon entering and exiting.
+        - 'cmd_line': (str) command-line ran for this profile.
+    """
+    func_ts = {}
+    mem_usage = []
+    timestamp = []
+    cmd_line = None
     f = open(filename, "r")
     for l in f:
-        fields = l.split()
-        if fields[0] == "MEM":
+        field, value = l.split(' ', 1)
+        if field == "MEM":
             # mem, timestamp
-            mdata.append((fields[1], fields[2]))
-
-        elif fields[0] == "FUNC":
-            f_name, mem_start, start, mem_end, end = fields[1:]
-            ts = ret.get(f_name, [])
-            ts.append([float(start), float(end), float(mem_start), float(mem_end)])
-            ret[f_name] = ts
-
+            values = value.split(' ')
+            mem_usage.append(float(values[0]))
+            timestamp.append(float(values[1]))
+
+        elif field == "FUNC":
+            values = value.split(' ')
+            f_name, mem_start, start, mem_end, end = values[:5]
+            ts = func_ts.get(f_name, [])
+            ts.append([float(start), float(end),
+                       float(mem_start), float(mem_end)])
+            func_ts[f_name] = ts
+
+        elif field == "CMDLINE":
+            cmd_line = value
         else:
             pass
     f.close()
 
-    mdata = np.asarray(mdata,
-                       dtype=[("mem", np.float), ("timestamp", np.float)])
-    return mdata, ret
+    return {"mem_usage": mem_usage, "timestamp": timestamp,
+            "func_timestamp": func_ts, 'filename': filename,
+            'cmd_line': cmd_line}
 
 
 
@@ -293,17 +321,41 @@ def plot_file(filename, index=0, timestamps=True):
     except ImportError:
         print("matplotlib is needed for plotting.")
         sys.exit(1)
+    import numpy as np  # pylab requires numpy anyway
+    mprofile = read_mprofile_file(filename)
+
+    if len(mprofile['timestamp']) == 0:
+        print('** No memory usage values have been found in the profile '
+              'file.**\nFile path: {0}\n'
+              'File may be empty or invalid.\n'
+              'It can be deleted with "mprof rm {0}"'.format(
+              mprofile['filename']))
+        sys.exit(0)
+
+    # Merge function timestamps and memory usage together
+    ts = mprofile['func_timestamp']
+    t = mprofile['timestamp']
+    mem = mprofile['mem_usage']
+
+    if len(ts) > 0:
+        for values in ts.itervalues():
+            for v in values:
+                t.extend(v[:2])
+                mem.extend(v[2:4])
+
+    mem = np.asarray(mem)
+    t = np.asarray(t)
+    ind = t.argsort()
+    mem = mem[ind]
+    t = t[ind]
+
+    # Plot curves
+    global_start = float(t[0])
+    t = t - global_start
 
-    mdata, ts = read_mprofile_file(filename)
-
-    global_start = float(mdata["timestamp"][0])
-
-    mem = mdata["mem"]
     max_mem = mem.max()
     max_mem_ind = mem.argmax()
 
-    t = mdata["timestamp"] - global_start
-
     all_colors=("c", "y", "g", "r", "b")
     mem_line_colors=('k', "b", "r")
     mem_line_label = time.strftime("%d / %m / %Y - start at %H:%M:%S",
@@ -334,6 +386,8 @@ def plot_file(filename, index=0, timestamps=True):
                   colors="r", linestyles="--")
         pl.vlines(t[max_mem_ind], bottom, top,
                   colors="r", linestyles="--")
+    return mprofile
+
 
 def plot_action():
     try:
@@ -342,18 +396,29 @@ def plot_action():
         print("matplotlib is needed for plotting.")
         sys.exit(1)
 
+    parser = OptionParser(version=mp.__version__)
+    parser.disable_interspersed_args()
+    parser.add_option("--title", "-t", dest="title", default=None,
+                      type="str", action="store",
+                      help="String shown as plot title")
+    parser.add_option("--no-function-ts", "-n", dest="no_timestamps",
+                      default=False, action="store_true",
+                      help="Do not display function timestamps on plot.")
+    (options, args) = parser.parse_args()
+
     profiles = glob.glob("mprofile_??????????????.dat")
     profiles.sort()
 
-    if len(sys.argv) == 1:
+    if len(args) == 0:
         if len(profiles) == 0:
             print("No input file found. \nThis program looks for "
-                  "mprofile_*.dat files, generated by the mprofile command.")
+                  "mprofile_*.dat files, generated by the "
+                  "'mprof run' command.")
             sys.exit(-1)
         filenames = [profiles[-1]]
     else:
         filenames = []
-        for arg in sys.argv[2:]:
+        for arg in args:
             if osp.exists(arg):
                 if not arg in filenames:
                     filenames.append(arg)
@@ -366,15 +431,21 @@ def plot_action():
                     filenames.append(profiles[n])
 
     pl.figure(figsize=(14, 6), dpi=90)
-    if len(filenames) > 1:
+    if len(filenames) > 1 or options.no_timestamps:
         timestamps = False
     else:
         timestamps = True
     for n, filename in enumerate(filenames):
-        plot_file(filename, index=n, timestamps=timestamps)
+        mprofile = plot_file(filename, index=n, timestamps=timestamps)
     pl.xlabel("time (in seconds)")
     pl.ylabel("memory used (in MiB)")
 
+    if options.title is None and len(filenames) == 1:
+        pl.title(mprofile['cmd_line'])
+    else:
+        if options.title is not None:
+            pl.title(options.title)
+
     ax = pl.gca()
     box = ax.get_position()
     ax.set_position([0.07, 0.1,
@@ -384,7 +455,8 @@ def plot_action():
     pl.show()
 
 if __name__ == "__main__":
-    # Workaround for optparse limitation: insert -- before first negative number found.
+    # Workaround for optparse limitation: insert -- before first negative
+    # number found.
     negint = re.compile("-[0-9]+")
     for n, arg in enumerate(sys.argv):
         if negint.match(arg):