GBench plot tool improvement

Reviewers: teon.banek

Reviewed By: teon.banek

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D760

tools/plot_gbench_json

@@ -11,13 +11,12 @@ Does a few nice things for you:
        This is currently implemented to work well with template based
        benchmarks, it might required mods.
     3. Automatically detects the need for log-scale on both axes.
+    4. Displaying plots or saving them all to a folder.
 
 Missing features:
     1. Proper support for benchmarks that use two arguments.
     2. Proper handling for all types of benchmark structures, name parsing
        in this implementation is made for template-based benches.
-    3. Plotting to image files. This implementation plots to the GUI (you can
-       save images there).
 
 Usage:
     # Generate benchmark data in json format using:
@@ -26,19 +25,36 @@ Usage:
     > ./plot_bench_json data.json
 
 Alternatively you can route stuff and avoid using an intermediary file:
-    sh > ./my_bench --benchmark_out_format=json --benchmark_out=/dev/stderr 2>&1 >/dev/null | grep "^[{} ]" | plot_gbench_json
+    sh > ./my_bench --benchmark_out_format=json
+         --benchmark_out=/dev/stderr 2>&1 >/dev/null
+         | grep "^[{} ]" | plot_gbench_json
 
 Maybe there is a nicer way to route it?
 """
 
+import sys
 import re
-import fileinput
+import os
 import json
+from argparse import ArgumentParser
 from collections import defaultdict
 
 from matplotlib import pyplot as plt
 
 
+def parse_args():
+    argp = ArgumentParser(description=__doc__)
+    argp.add_argument("input_file", nargs="?",
+                      help="Path to file with JSON data. If not provided data "
+                      "is read from stdin")
+    argp.add_argument("--output-path", required=False,
+                      help="Path to a folder where the plots should be saved. "
+                      " If not provided the plots are displayed in the GUI.")
+    argp.add_argument("--output-type", required=False, default="png",
+                      help="If saving plot files use this extension.")
+    return argp.parse_args()
+
+
 def convert_num(string):
     """
     Converts stuff like "100" and "3k" to numbers.
@@ -75,7 +91,12 @@ def is_exponential_growth(numbers):
 
 
 def main():
-    data = json.loads("".join(fileinput.input()).strip())
+    args = parse_args()
+    if args.input_file:
+        with open(args.input_file) as f:
+            data = json.load(f)
+    else:
+        data = json.load(sys.stdin)
 
     # structure: {bench_name: [(x, y, time_unit), ...]
     benchmarks = defaultdict(list)
@@ -121,7 +142,13 @@ def main():
         plt.title(group_name)
         plt.legend()
         plt.grid()
-        plt.show()
+        if args.output_path:
+            if not os.path.exists(args.output_path):
+                os.makedirs(args.output_path, exist_ok=True)
+            plt.savefig(os.path.join(
+                args.output_path, group_name + "." + args.output_type))
+        else:
+            plt.show()
 
 
 if __name__ == "__main__":