Merge pull request #302 from william-silversmith/master

`mprof peak` displays maximum memory usage of .dat files.

Author: fabianp

mprof.py

@@ -16,7 +16,7 @@
 import importlib
 import memory_profiler as mp
 
-ALL_ACTIONS = ("run", "rm", "clean", "list", "plot", "attach")
+ALL_ACTIONS = ("run", "rm", "clean", "list", "plot", "attach", "peak")
 help_msg = """
 Available commands:
 
@@ -26,6 +26,7 @@
     clean    clean the current directory from files created by mprof
     list     display existing profiles, with indices
     plot     plot memory consumption generated by mprof run
+    peak     print the maximum memory used by an mprof run
 
 Type mprof <command> --help for usage help on a specific command.
 For example, mprof plot --help will list all plotting options.
@@ -73,7 +74,7 @@ def get_profile_filenames(args):
     profiles = glob.glob("mprofile_??????????????.dat")
     profiles.sort()
 
-    if args is "all":
+    if args == "all":
         filenames = copy.copy(profiles)
     else:
         filenames = []
@@ -794,33 +795,7 @@ def xlim_type(value):
         sys.exit(1)
     pl.ioff()
 
-    profiles = glob.glob("mprofile_??????????????.dat")
-    profiles.sort()
-
-    if len(args.profiles) == 0:
-        if len(profiles) == 0:
-            print("No input file found. \nThis program looks for "
-                  "mprofile_*.dat files, generated by the "
-                  "'mprof run' command.")
-            sys.exit(-1)
-        print("Using last profile data.")
-        filenames = [profiles[-1]]
-    else:
-        filenames = []
-        for prof in args.profiles:
-            if osp.exists(prof):
-                if not prof in filenames:
-                    filenames.append(prof)
-            else:
-                try:
-                    n = int(prof)
-                    if not profiles[n] in filenames:
-                        filenames.append(profiles[n])
-                except ValueError:
-                    print("Input file not found: " + prof)
-    if not len(filenames):
-        print("No files found from given input.")
-        sys.exit(-1)
+    filenames = get_profiles(args)
 
     fig = pl.figure(figsize=(14, 6), dpi=90)
     if not args.flame_mode:
@@ -860,6 +835,55 @@ def xlim_type(value):
     else:
         pl.show()
 
+def peak_action():
+    desc = """Prints the peak memory used in data file `file.dat` generated
+using `mprof run`. If no .dat file is given, it will take the most recent
+such file in the current directory."""
+    parser = ArgumentParser(usage="mprof peak [options] [file.dat]", description=desc)
+    parser.add_argument("profiles", nargs="*",
+                    help="profiles made by mprof run")
+    args = parser.parse_args()
+    filenames = get_profiles(args)
+
+    for filename in filenames:
+        prof = read_mprofile_file(filename)
+        print("{}\t{:.3f} MiB".format(prof["filename"], max(prof["mem_usage"])))
+        for child, values in prof["children"].items():
+            child_peak = max([ mem_ts[0] for mem_ts in values ])
+            print("  Child {}\t\t\t{:.3f} MiB".format(child, child_peak))
+        
+
+def get_profiles(args):
+    profiles = glob.glob("mprofile_??????????????.dat")
+    profiles.sort()
+
+    if len(args.profiles) == 0:
+        if len(profiles) == 0:
+            print("No input file found. \nThis program looks for "
+                  "mprofile_*.dat files, generated by the "
+                  "'mprof run' command.")
+            sys.exit(-1)
+        print("Using last profile data.")
+        filenames = [profiles[-1]]
+    else:
+        filenames = []
+        for prof in args.profiles:
+            if osp.exists(prof):
+                if not prof in filenames:
+                    filenames.append(prof)
+            else:
+                try:
+                    n = int(prof)
+                    if not profiles[n] in filenames:
+                        filenames.append(profiles[n])
+                except ValueError:
+                    print("Input file not found: " + prof)
+    if not len(filenames):
+        print("No files found from given input.")
+        sys.exit(-1)
+
+    return filenames
+
 def main():
     # Workaround for optparse limitation: insert -- before first negative
     # number found.
@@ -873,7 +897,8 @@ def main():
                "list": list_action,
                "run": run_action,
                "attach": attach_action,
-               "plot": plot_action}
+               "plot": plot_action,
+               "peak": peak_action}
     actions[get_action()]()
 
 if __name__ == "__main__":