Merge pull request #769 from jdufresne/argparse

Switch from using optparse to argparse for command line arguments

locust/main.py

@@ -6,7 +6,7 @@
 import socket
 import sys
 import time
-from optparse import OptionParser
+import argparse
 
 import gevent
 
@@ -26,276 +26,222 @@
 
 def parse_options():
     """
-    Handle command-line options with optparse.OptionParser.
+    Handle command-line options with argparse.ArgumentParser.
 
     Return list of arguments, largely for use in `parse_arguments`.
     """
 
     # Initialize
-    parser = OptionParser(usage="locust [options] [LocustClass [LocustClass2 ... ]]")
+    parser = argparse.ArgumentParser()
 
-    parser.add_option(
+    parser.add_argument(
         '-H', '--host',
-        dest="host",
-        default=None,
         help="Host to load test in the following format: http://10.21.32.33"
     )
 
-    parser.add_option(
+    parser.add_argument(
         '--web-host',
-        dest="web_host",
         default="",
         help="Host to bind the web interface to. Defaults to '' (all interfaces)"
     )
 
-    parser.add_option(
+    parser.add_argument(
         '-P', '--port', '--web-port',
-        type="int",
-        dest="port",
+        type=int,
         default=8089,
         help="Port on which to run web host"
     )
 
-    parser.add_option(
+    parser.add_argument(
         '-f', '--locustfile',
-        dest='locustfile',
         default='locustfile',
         help="Python module file to import, e.g. '../other.py'. Default: locustfile"
     )
 
     # A file that contains the current request stats.
-    parser.add_option(
+    parser.add_argument(
         '--csv', '--csv-base-name',
-        action='store',
-        type='str',
         dest='csvfilebase',
-        default=None,
         help="Store current request stats to files in CSV format.",
     )
 
     # if locust should be run in distributed mode as master
-    parser.add_option(
+    parser.add_argument(
         '--master',
         action='store_true',
-        dest='master',
-        default=False,
         help="Set locust to run in distributed mode with this process as master"
     )
 
     # if locust should be run in distributed mode as slave
-    parser.add_option(
+    parser.add_argument(
         '--slave',
         action='store_true',
-        dest='slave',
-        default=False,
         help="Set locust to run in distributed mode with this process as slave"
     )
 
     # master host options
-    parser.add_option(
+    parser.add_argument(
         '--master-host',
-        action='store',
-        type='str',
-        dest='master_host',
         default="127.0.0.1",
         help="Host or IP address of locust master for distributed load testing. Only used when running with --slave. Defaults to 127.0.0.1."
     )
 
-    parser.add_option(
+    parser.add_argument(
         '--master-port',
-        action='store',
-        type='int',
-        dest='master_port',
+        type=int,
         default=5557,
         help="The port to connect to that is used by the locust master for distributed load testing. Only used when running with --slave. Defaults to 5557. Note that slaves will also connect to the master node on this port + 1."
     )
 
-    parser.add_option(
+    parser.add_argument(
         '--master-bind-host',
-        action='store',
-        type='str',
-        dest='master_bind_host',
         default="*",
         help="Interfaces (hostname, ip) that locust master should bind to. Only used when running with --master. Defaults to * (all available interfaces)."
     )
 
-    parser.add_option(
+    parser.add_argument(
         '--master-bind-port',
-        action='store',
-        type='int',
-        dest='master_bind_port',
+        type=int,
         default=5557,
         help="Port that locust master should bind to. Only used when running with --master. Defaults to 5557. Note that Locust will also use this port + 1, so by default the master node will bind to 5557 and 5558."
     )
 
-    parser.add_option(
+    parser.add_argument(
         '--heartbeat-liveness',
-        action='store',
-        type='int',
-        dest='heartbeat_liveness',
+        type=int,
         default=3,
         help="set number of seconds before failed heartbeat from slave"
     )
 
-    parser.add_option(
+    parser.add_argument(
         '--heartbeat-interval',
-        action='store',
-        type='int',
-        dest='heartbeat_interval',
+        type=int,
         default=1,
         help="set number of seconds delay between slave heartbeats to master"
     )
 
-    parser.add_option(
+    parser.add_argument(
         '--expect-slaves',
-        action='store',
-        type='int',
-        dest='expect_slaves',
+        type=int,
         default=1,
         help="How many slaves master should expect to connect before starting the test (only when --no-web used)."
     )
 
     # if we should print stats in the console
-    parser.add_option(
+    parser.add_argument(
         '--no-web',
         action='store_true',
-        dest='no_web',
-        default=False,
         help="Disable the web interface, and instead start running the test immediately. Requires -c and -r to be specified."
     )
 
     # Number of clients
-    parser.add_option(
+    parser.add_argument(
         '-c', '--clients',
-        action='store',
-        type='int',
+        type=int,
         dest='num_clients',
         default=1,
         help="Number of concurrent Locust users. Only used together with --no-web"
     )
 
     # Client hatch rate
-    parser.add_option(
+    parser.add_argument(
         '-r', '--hatch-rate',
-        action='store',
-        type='float',
-        dest='hatch_rate',
+        type=float,
         default=1,
         help="The rate per second in which clients are spawned. Only used together with --no-web"
     )
 
     # Time limit of the test run
-    parser.add_option(
+    parser.add_argument(
         '-t', '--run-time',
-        action='store',
-        type='str',
-        dest='run_time',
-        default=None,
         help="Stop after the specified amount of time, e.g. (300s, 20m, 3h, 1h30m, etc.). Only used together with --no-web"
     )
 
     # log level
-    parser.add_option(
+    parser.add_argument(
         '--loglevel', '-L',
-        action='store',
-        type='str',
-        dest='loglevel',
         default='INFO',
         help="Choose between DEBUG/INFO/WARNING/ERROR/CRITICAL. Default is INFO.",
     )
 
     # log file
-    parser.add_option(
+    parser.add_argument(
         '--logfile',
-        action='store',
-        type='str',
-        dest='logfile',
-        default=None,
         help="Path to log file. If not set, log will go to stdout/stderr",
     )
 
     # if we should print stats in the console
-    parser.add_option(
+    parser.add_argument(
         '--print-stats',
         action='store_true',
-        dest='print_stats',
-        default=False,
         help="Print stats in the console"
     )
 
     # only print summary stats
-    parser.add_option(
+    parser.add_argument(
        '--only-summary',
        action='store_true',
-       dest='only_summary',
-       default=False,
        help='Only print the summary stats'
     )
 
-    parser.add_option(
+    parser.add_argument(
         '--no-reset-stats',
         action='store_true',
         help="[DEPRECATED] Do not reset statistics once hatching has been completed. This is now the default behavior. See --reset-stats to disable",
     )
 
-    parser.add_option(
+    parser.add_argument(
         '--reset-stats',
         action='store_true',
-        dest='reset_stats',
-        default=False,
         help="Reset statistics once hatching has been completed. Should be set on both master and slaves when running in distributed mode",
     )
 
     # List locust commands found in loaded locust files/source files
-    parser.add_option(
+    parser.add_argument(
         '-l', '--list',
         action='store_true',
         dest='list_commands',
-        default=False,
         help="Show list of possible locust classes and exit"
     )
 
     # Display ratio table of all tasks
-    parser.add_option(
+    parser.add_argument(
         '--show-task-ratio',
         action='store_true',
-        dest='show_task_ratio',
-        default=False,
         help="print table of the locust classes' task execution ratio"
     )
     # Display ratio table of all tasks in JSON format
-    parser.add_option(
+    parser.add_argument(
         '--show-task-ratio-json',
         action='store_true',
-        dest='show_task_ratio_json',
-        default=False,
         help="print json data of the locust classes' task execution ratio"
     )
 
     # Version number (optparse gives you --version but we have to do it
     # ourselves to get -V too. sigh)
-    parser.add_option(
+    parser.add_argument(
         '-V', '--version',
-        action='store_true',
-        dest='show_version',
-        default=False,
-        help="show program's version number and exit"
+        action='version',
+        version='%(prog)s {}'.format(version),
     )
 
     # set the exit code to post on errors
-    parser.add_option(
+    parser.add_argument(
         '--exit-code-on-error',
-        action='store',
-        type="int",
-        dest='exit_code_on_error',
+        type=int,
         default=1,
         help="sets the exit code to post on error"
     )
 
+    parser.add_argument(
+        'locust_classes',
+        nargs='*',
+        metavar='LocustClass',
+    )
+
     # Finalize
-    # Return three-tuple of parser + the output from parse_args (opt obj, args)
-    opts, args = parser.parse_args()
-    return parser, opts, args
+    # Return two-tuple of parser + the output from parse_args
+    return parser, parser.parse_args()
 
 
 def _is_package(path):
@@ -413,15 +359,11 @@ def __import_locustfile__(filename, path):
     return imported.__doc__, locusts
 
 def main():
-    parser, options, arguments = parse_options()
+    parser, options = parse_options()
 
     # setup logging
     setup_logging(options.loglevel, options.logfile)
     logger = logging.getLogger(__name__)
-
-    if options.show_version:
-        print("Locust %s" % (version,))
-        sys.exit(0)
 
     locustfile = find_locustfile(options.locustfile)
 
@@ -446,13 +388,13 @@ def main():
         sys.exit(1)
 
     # make sure specified Locust exists
-    if arguments:
-        missing = set(arguments) - set(locusts.keys())
+    if options.locust_classes:
+        missing = set(options.locust_classes) - set(locusts.keys())
         if missing:
             logger.error("Unknown Locust(s): %s\n" % (", ".join(missing)))
             sys.exit(1)
         else:
-            names = set(arguments) & set(locusts.keys())
+            names = set(options.locust_classes) & set(locusts.keys())
             locust_classes = [locusts[n] for n in names]
     else:
         # list() call is needed to consume the dict_view object in Python 3

locust/test/test_parser.py

@@ -5,22 +5,22 @@
 
 class TestParser(unittest.TestCase):
     def setUp(self):
-        self.parser, _, _ = parse_options()
+        self.parser, _ = parse_options()
 
     def test_default(self):
-        opts, _ = self.parser.parse_args([])
+        opts = self.parser.parse_args([])
         self.assertEqual(opts.reset_stats, False)
 
     def test_reset_stats(self):
         args = [
             "--reset-stats"
         ]
-        opts, _ = self.parser.parse_args(args)
+        opts = self.parser.parse_args(args)
         self.assertEqual(opts.reset_stats, True)
 
     def test_should_accept_legacy_no_reset_stats(self):
         args = [
             "--no-reset-stats"
         ]
-        opts, _ = self.parser.parse_args(args)
+        opts = self.parser.parse_args(args)
         self.assertEqual(opts.reset_stats, False)