From 783a421a3c1e10205f9668238a1271767a4083c0 Mon Sep 17 00:00:00 2001
From: Calen Pennington <cale@edx.org>
Date: Mon, 9 Mar 2015 14:18:30 -0400
Subject: [PATCH] Add distribution and request csvs to commandline usage.

Original-Author: Kevin Peter Savage (https://github.com/kevinpetersavage)
---
 locust/main.py | 50 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/locust/main.py b/locust/main.py
index a7b41a6d53..5d8a0162b6 100644
--- a/locust/main.py
+++ b/locust/main.py
@@ -8,6 +8,7 @@
 import inspect
 import logging
 import socket
+import time
 from optparse import OptionParser
 
 import web
@@ -228,6 +229,23 @@ def parse_options():
         help="show program's version number and exit"
     )
 
+    # Writes output statistics as csv
+    parser.add_option(
+        '--write-stats',
+        action='store_true',
+        dest='write_stats',
+        default=None,
+        help="writes output statistics as csv files"
+    )
+
+    parser.add_option(
+        '--stats-prefix',
+        action='store',
+        dest='stats_prefix',
+        default=None,
+        help="path prefix to use when writing stats files"
+    )
+
     # Finalize
     # Return three-tuple of parser + the output from parse_args (opt obj, args)
     opts, args = parser.parse_args()
@@ -290,6 +308,28 @@ def is_locust(tup):
     )
 
 
+def write_distribution_stats_csv(file_prefix, timestamp=None):
+    if timestamp is None:
+        timestamp = time.time()
+
+    file_name = "{0}distribution_{1}.csv".format(file_prefix, timestamp)
+    console_logger.info("Writing distribution stats to %s", file_name)
+    with open(file_name, 'wb') as file:
+        data = runners.locust_runner.stats.get_percentile_dataset(include_empty=True)
+        file.write(data.csv)
+
+
+def write_stats_csv(file_prefix, timestamp=None):
+    if timestamp is None:
+        timestamp = time.time()
+
+    file_name = "{0}stats_{1}.csv".format(file_prefix, timestamp)
+    console_logger.info("Writing request stats to %s", file_name)
+    with open(file_name, 'wb') as file:
+        data = runners.locust_runner.stats.get_request_stats_dataset()
+        file.write(data.csv)
+
+
 def load_locustfile(path):
     """
     Import given locustfile path and return (docstring, callables).
@@ -442,6 +482,16 @@ def sig_term_handler():
         code = 0
         if len(runners.locust_runner.errors):
             code = 1
+
+        if options.write_stats:
+            if options.stats_prefix is not None:
+                file_prefix = os.path.expanduser(options.stats_prefix) + '_'
+            else:
+                file_prefix = ''
+            timestamp = time.time()
+            write_distribution_stats_csv(file_prefix, timestamp)
+            write_stats_csv(file_prefix, timestamp)
+
         shutdown(code=code)
     except KeyboardInterrupt as e:
         shutdown(0)