Refactored command line parser to work better with subcommands

Environment variable support is now part of the command line tool only,
not the library itself.

Author: Sam Clements

riemann_client/command.py

@@ -4,6 +4,7 @@
 
 import argparse
 import json
+import os
 
 import riemann_client.client
 import riemann_client.transport
@@ -13,64 +14,41 @@
     'tcp': riemann_client.transport.TCPTransport
 }
 
+parser = argparse.ArgumentParser(add_help=False, description=(
+    "Uses the RIEMANN_HOST and RIEMANN_PORT environment variables "
+    "if no host and port are given. If they are not set, the transports "
+    "will use a default host and port of localhost:5555."))
 
-def wide_formatter(*args, **kwargs):
-    kwargs.setdefault('max_help_position', 32)
-    kwargs.setdefault('width', 96)
-    return argparse.HelpFormatter(*args, **kwargs)
+parser.add_argument(
+    '-h', '--help', action='help',
+    help="show this help message and exit")
 
-parser = argparse.ArgumentParser(
-    formatter_class=wide_formatter, description=(
-        "Uses the RIEMANN_HOST and RIEMANN_PORT environment variables "
-        "if no host and port are given. If they are not set, the transports "
-        "will use a default host and port of localhost:5555."))
 parser.add_argument(
     '-v', '--version', action='version',
     version='python-riemann-client v{version} by {author}'.format(
         version=riemann_client.__version__,
         author=riemann_client.__author__),
     help="Show this program's version and exit")
+
 parser.add_argument(
-    'host', type=str, nargs='?', default=None,
-    help="The hostname of a Riemann server (default: localhost)")
+    '-H', '--host', type=str,
+    default=os.environ.get('RIEMANN_HOST', 'localhost'),
+    help="The hostname of a Riemann server (environ: %(default)s)")
+
 parser.add_argument(
-    'port', type=int, nargs='?', default=None,
-    help="The port to connect to the Riemann server on (default: 5555)")
+    '-P', '--port', type=int,
+    default=os.environ.get('RIEMANN_PORT', 5555),
+    help="The port to connect to the Riemann server on (environ: %(default)s)")
+
 parser.add_argument(
     '-t', '--transport', choices=TRANSPORT_CLASSES.keys(), default='tcp',
     help="The transport to use (default: %(default)s)")
 
 subparsers = parser.add_subparsers(dest='subparser')
 
-send = subparsers.add_parser(
-    'send', formatter_class=wide_formatter,
-    help='Send an event to Riemann')
-send.add_argument(
-    '-p', '--print', action='store_true', dest='print_message',
-    help="Print the message that is sent to Riemann")
-
-for arg_names, arg_type, arg_help in [
-        (['-u', '--time'], int, "Unix timestamp"),
-        (['-S', '--state'], str, "State"),
-        (['-e', '--event-host'], str, "Host"),
-        (['-D', '--description'], str, "Description"),
-        (['-s', '--service'], str, "Service"),
-        (['-T', '--tags'], list, "Tags"),
-        (['-l', '--ttl'], int, "Time to live"),
-        (['-m', '--metric'], float, "Value")]:
-    send.add_argument(
-        *arg_names, metavar=arg_type.__name__, type=arg_type, help=arg_help)
-
-query = subparsers.add_parser('query', help='Query the Riemann index')
-query.add_argument(
-    'query', type=str,
-    help="The query to send")
-query.add_argument(
-    '-pp', '--pretty-print', action='store_true',
-    help="Pretty print output")
-
 
-def send(args, client):
+def send_function(args, client):
+    """Sends a single command to Riemann"""
     event = client.event(
         time=args.time,
         state=args.state,
@@ -84,13 +62,44 @@ def send(args, client):
     if args.print_message:
         print(str(event).strip())
 
+send = subparsers.add_parser('send', help='Send an event to Riemann')
+
+send.add_argument(
+    '-p', '--print', action='store_true', dest='print_message',
+    help="Print the message that is sent to Riemann")
+
+send_arguments = [
+    (['-u', '--time'], int, "Unix timestamp"),
+    (['-S', '--state'], str, "State"),
+    (['-e', '--event-host'], str, "Host"),
+    (['-D', '--description'], str, "Description"),
+    (['-s', '--service'], str, "Service"),
+    (['-T', '--tags'], list, "Tags"),
+    (['-l', '--ttl'], int, "Time to live"),
+    (['-m', '--metric'], float, "Value")]
+
+for flags, atype, ahelp in send_arguments:
+    send.add_argument(*flags, metavar=atype.__name__, type=atype, help=ahelp)
+
+send.set_defaults(function=send_function)
 
-def query(args, client):
+
+def query_function(args, client):
+    """Queries the Riemann index and outputs the returned events as JSON
+
+    UDP is not supported when querying"""
     print(json.dumps(client.query(args.query), sort_keys=True, indent=2))
 
+query = subparsers.add_parser('query', help='Query the Riemann index')
+query.add_argument('query', type=str, help="The query to send")
+query.set_defaults(function=query_function)
+
 
 def main():
     args = parser.parse_args()
-    transport = TRANSPORT_CLASSES[args.transport](args.host, args.port)
+
+    transport_class = TRANSPORT_CLASSES[args.transport]
+    transport = transport_class(args.host, args.port)
+
     with riemann_client.client.Client(transport=transport) as client:
-        {'send': send, 'query': query}[args.subparser](args, client)
+        args.function(args, client)

riemann_client/transport.py

@@ -3,7 +3,6 @@
 from __future__ import absolute_import
 
 import abc
-import os
 import socket
 import struct
 
@@ -18,9 +17,9 @@ class RiemannError(Exception):
 class Transport(object):
     __metaclass__ = abc.ABCMeta
 
-    def __init__(self, host=None, port=None):
-        self.host = host or os.environ.get('RIEMANN_HOST', 'localhost')
-        self.port = port or os.environ.get('RIEMANN_PORT', 5555)
+    def __init__(self, host='localhost', port=5555):
+        self.host = host
+        self.port = port
 
     def __enter__(self):
         self.connect()