Merge pull request #187 from Yelp/performance-check

Adding baseline functionality for benchmark script

Author: domanchi

Diff for: detect_secrets/util.py

@@ -0,0 +1,10 @@
+import os
+
+
+def get_root_directory():
+    return os.path.realpath(
+        os.path.join(
+            os.path.dirname(__file__),
+            '../',
+        ),
+    )

Diff for: scripts/benchmark.py

@@ -10,7 +10,10 @@
 
 from monotonic import monotonic
 
+from detect_secrets.core.color import AnsiColor
+from detect_secrets.core.color import colorize
 from detect_secrets.core.usage import PluginOptions
+from detect_secrets.util import get_root_directory
 
 
 def main():
@@ -92,7 +95,7 @@ def get_arguments():
     parser.add_argument(
         '--harakiri',
         default=5,
-        type=float,
+        type=assert_positive(float),
         help=(
             'Specifies an upper bound for the number of seconds to wait '
             'per execution.'
@@ -102,40 +105,59 @@ def get_arguments():
         '-n',
         '--num-iterations',
         default=1,
-        type=assert_positive_integer,
+        type=assert_positive(int),
         help=(
             'Specifies the number of times to run the test. '
             'Results will be averaged over this value.'
         ),
     )
+    parser.add_argument(
+        '--baseline',
+        type=assert_valid_file,
+        help=(
+            'If provided, will compare performance with provided baseline. '
+            'Assumes pretty output (otherwise, you can do the comparison '
+            'yourself).'
+        ),
+    )
 
     args = parser.parse_args()
     if not args.filenames:
-        args.filenames = [
-            os.path.realpath(
-                os.path.join(
-                    os.path.dirname(__file__),
-                    '../',
-                ),
-            ),
-        ]
+        if args.baseline:
+            args.filenames = args.baseline['filenames']
+        else:
+            args.filenames = [get_root_directory()]
 
     if not args.plugin:
         args.plugin = plugins
 
     return args
 
 
-def assert_positive_integer(string):
-    value = int(string)
-    if value <= 0:
+def assert_positive(type):
+    def wrapped(string):
+        value = type(string)
+        if value <= 0:
+            raise argparse.ArgumentTypeError(
+                '{} must be a positive {}.'.format(
+                    string,
+                    type.__name__,
+                ),
+            )
+
+        return value
+
+    return wrapped
+
+
+def assert_valid_file(string):
+    if not os.path.isfile(string):
         raise argparse.ArgumentTypeError(
-            '{} must be a positive integer.'.format(
-                string,
-            ),
+            '{} must be a valid file.'.format(string),
         )
 
-    return value
+    with open(string) as f:
+        return json.load(f)
 
 
 def time_execution(filenames, timeout, num_iterations=1, flags=None):
@@ -166,39 +188,125 @@ def time_execution(filenames, timeout, num_iterations=1, flags=None):
     if result == timeout:
         return None
 
-    return statistics.mean(scores)
+    return round(statistics.mean(scores), 5)
 
 
 def print_output(timings, args):
     """
     :type timings: dict
     :type args: Namespace
     """
-    if not args.pretty:
-        print(json.dumps(timings))
+    if not args.pretty and not args.baseline:
+        print(
+            json.dumps({
+                'filenames': args.filenames,
+                'timings': timings,
+            }),
+        )
         return
 
     # Print header
-    print('-' * 42)
-    print('{:<20s}{:>20s}'.format('plugin', 'time'))
-    print('-' * 42)
+    baseline = args.baseline['timings'] if args.baseline else {}
+    if not baseline:
+        print('-' * 45)
+        print('{:<25s}{:>15s}'.format('plugin', 'time'))
+        print('-' * 45)
+    else:
+        print('-' * 60)
+        print('{:<25s}{:>11s}{:>22s}'.format('plugin', 'time', 'change'))
+        print('-' * 60)
 
+    # Print content
     if 'all-plugins' in timings:
-        print_line('all-plugins', timings['all-plugins'])
+        print_line(
+            'All Plugins',
+            time=timings['all-plugins'],
+            baseline=_get_baseline_value(baseline, 'all-plugins'),
+            timeout=args.harakiri,
+        )
         del timings['all-plugins']
 
     for key in sorted(timings):
-        print_line(key, timings[key])
-    print('-' * 42)
+        print_line(
+            key,
+            time=timings[key],
+            baseline=_get_baseline_value(baseline, key),
+            timeout=args.harakiri,
+        )
 
+    # Print footer line
+    if not args.baseline:
+        print('-' * 45)
+    else:
+        print('-' * 60)
+
+
+def _get_baseline_value(baseline, key):
+    """
+    We need to distinguish between no baseline mode (which should return
+    None as a value), baseline mode with exceeded timeout (which is stored
+    as None, but should return 0).
+    """
+    if key in baseline:
+        return 0 if baseline[key] is None else baseline[key]
+
+
+def print_line(name, time, baseline, timeout):
+    """
+    :type name: str
+
+    :type time: float
+    :param time: seconds it took to execute
+
+    :type baseline: float
+    :param baseline: expected seconds to execute
 
-def print_line(name, time):
+    :type timeout: float
+    :param timeout: used to calculate difference when either current
+        execution or baseline execution exceeds timeout.
+    """
     if not time:
-        time = 'Timeout exceeded!'
+        time_string = 'Timeout exceeded!'
     else:
-        time = '{}s'.format(str(time))
+        time_string = '{}s'.format(str(time))
+
+    if baseline is not None:
+        if time and baseline:
+            difference = round(baseline - time, 2)
+        elif time:
+            # This handles the case when the baseline execution exceeds timeout
+            difference = round(timeout - time, 2)
+        elif baseline:
+            # This handles the case when this current execution exceeds timeout
+            difference = round(timeout - baseline, 2)
+        else:
+            # They both failed.
+            difference = 0
+
+        if difference > 0:
+            difference_string = colorize(
+                '▲  {}'.format(difference),
+                AnsiColor.LIGHT_GREEN,
+            )
+            difference_string = '{:>22s}'.format(difference_string)
+        elif difference < 0:
+            difference_string = colorize(
+                '▼ {}'.format(difference),
+                AnsiColor.RED,
+            )
+            difference_string = '{:>22s}'.format(difference_string)
+        else:
+            difference_string = '{:>10s}'.format('-')
 
-    print('{:<20s}{:>20s}'.format(name, time))
+        print(
+            '{:<25s}{:^20s}{}'.format(
+                name,
+                time_string,
+                difference_string,
+            ),
+        )
+    else:
+        print('{:<25s}{:>20s}'.format(name, time_string))
 
 
 if __name__ == '__main__':