Use click instead of argparse in tracker (#3599)

Author: Code0x58

heron/tools/tracker/src/python/BUILD

@@ -7,10 +7,11 @@ pex_library(
         exclude = ["main.py"],
     ),
     reqs = [
-        "protobuf==3.8.0",
-        "tornado==4.0.2",
+        "click==7.1.2",
         "javaobj-py3==0.4.1",
         "networkx==2.4",
+        "protobuf==3.8.0",
+        "tornado==4.0.2",
     ],
     deps = [
         "//heron/common/src/python:common-py",

heron/tools/tracker/src/python/main.py

@@ -19,7 +19,7 @@
 #  under the License.
 
 ''' main.py '''
-import argparse
+import logging
 import os
 import signal
 import sys
@@ -29,14 +29,16 @@
 from tornado.options import define
 from tornado.httpclient import AsyncHTTPClient
 
-import heron.tools.common.src.python.utils.config as common_config
-import heron.common.src.python.utils.log as log
+from heron.tools.common.src.python.utils import config as common_config
+from heron.common.src.python.utils import log
 from heron.tools.tracker.src.python import constants
 from heron.tools.tracker.src.python import handlers
 from heron.tools.tracker.src.python import utils
 from heron.tools.tracker.src.python.config import Config, STATEMGRS_KEY
 from heron.tools.tracker.src.python.tracker import Tracker
 
+import click
+
 Log = log.Log
 
 class Application(tornado.web.Application):
@@ -93,173 +95,110 @@ def __init__(self, config):
   def stop(self):
     self.tracker.stop_sync()
 
-# pylint: disable=protected-access
-class _HelpAction(argparse._HelpAction):
-  """ HelpAction """
-  def __call__(self, parser, namespace, values, option_string=None):
-    parser.print_help()
-
-    # retrieve subparsers from parser
-    subparsers_actions = [
-        action for action in parser._actions
-        if isinstance(action, argparse._SubParsersAction)]
-
-    # there will probably only be one subparser_action,
-    # but better save than sorry
-    for subparsers_action in subparsers_actions:
-      # get all subparsers and print help
-      for choice, subparser in list(subparsers_action.choices.items()):
-        print("Subparser '{}'".format(choice))
-        print(subparser.format_help())
-
-    parser.exit()
-
-# pylint: disable=bad-super-call
-class SubcommandHelpFormatter(argparse.RawDescriptionHelpFormatter):
-  """ Subcommand help formatter """
-  def _format_action(self, action):
-    parts = super(argparse.RawDescriptionHelpFormatter, self)._format_action(action)
-    if action.nargs == argparse.PARSER:
-      parts = "\n".join(parts.split("\n")[1:])
-    return parts
-
-
-def add_titles(parser):
-  """ add titles """
-  parser._positionals.title = "Required arguments"
-  parser._optionals.title = "Optional arguments"
-  return parser
-
-
-def add_arguments(parser):
-  """ add arguments """
-  default_config_file = os.path.join(
-      utils.get_heron_tracker_conf_dir(), constants.DEFAULT_CONFIG_FILE)
-
-  parser.add_argument(
-      '--config-file',
-      metavar='(a string; path to config file; default: "' + default_config_file + '")',
-      default=default_config_file)
-
-  parser.add_argument(
-      '--type',
-      metavar='(an string; type of state manager (zookeeper or file, etc.); example: ' \
-        + str(constants.DEFAULT_STATE_MANAGER_TYPE) + ')',
-      choices=["file", "zookeeper"])
-
-  parser.add_argument(
-      '--name',
-      metavar='(an string; name to be used for the state manager; example: ' \
-        + str(constants.DEFAULT_STATE_MANAGER_NAME) + ')')
-
-  parser.add_argument(
-      '--rootpath',
-      metavar='(an string; where all the states are stored; example: ' \
-        + str(constants.DEFAULT_STATE_MANAGER_ROOTPATH) + ')')
-
-  parser.add_argument(
-      '--tunnelhost',
-      metavar='(an string; if ssh tunneling needs to be established to connect to it; example: ' \
-        + str(constants.DEFAULT_STATE_MANAGER_TUNNELHOST) + ')')
-
-  parser.add_argument(
-      '--hostport',
-      metavar='(an string; only used to connect to zk, must be of the form \'host:port\';'\
-      ' example: ' + str(constants.DEFAULT_STATE_MANAGER_HOSTPORT) + ')')
-
-  parser.add_argument(
-      '--port',
-      metavar='(an integer; port to listen; default: ' + str(constants.DEFAULT_PORT) + ')',
-      type=int,
-      default=constants.DEFAULT_PORT)
 
-  parser.add_argument(
-      '--verbose',
-      action='store_true')
-
-  return parser
-
-def create_parsers():
-  """ create argument parser """
-  parser = argparse.ArgumentParser(
-      epilog='For detailed documentation, go to http://github.com/apache/incubator-heron',
-      usage="%(prog)s [options] [help]",
-      add_help=False)
-
-  parser = add_titles(parser)
-  parser = add_arguments(parser)
-
-  ya_parser = argparse.ArgumentParser(
-      parents=[parser],
-      formatter_class=SubcommandHelpFormatter,
-      add_help=False)
-
-  subparsers = ya_parser.add_subparsers(
-      title="Available commands")
-
-  help_parser = subparsers.add_parser(
-      'help',
-      help='Prints help',
-      add_help=False)
-
-  help_parser.set_defaults(help=True)
-
-  subparsers.add_parser(
-      'version',
-      help='Prints version',
-      add_help=True)
-
-  return parser, ya_parser
-
-def define_options(port, config_file):
+def define_options(port: int, config_file: str) -> None:
   """ define Tornado global variables """
   define("port", default=port)
   define("config_file", default=config_file)
 
-def create_tracker_config(namespace):
+
+def create_tracker_config(config_file: str, stmgr_override: dict) -> dict:
   # try to parse the config file if we find one
-  config_file = namespace["config_file"]
   config = utils.parse_config_file(config_file)
   if config is None:
-    Log.debug("Config file does not exists: %s" % config_file)
+    Log.debug(f"Config file does not exists: {config_file}")
     config = {STATEMGRS_KEY:[{}]}
 
-  # update the config if we have any flags
-  config_flags = ["type", "name", "rootpath", "tunnelhost", "hostport"]
-  config_to_update = config[STATEMGRS_KEY][0]
-  for flag in config_flags:
-    value = namespace.get(flag, None)
-    if value is not None:
-      config_to_update[flag] = value
-
+  # update non-null options
+  config[STATEMGRS_KEY][0].update(
+      (k, v)
+      for k, v in stmgr_override.items()
+      if v is not None
+  )
   return config
 
-def main():
-  """ main """
-  # create the parser and parse the arguments
-  (parser, _) = create_parsers()
-  (args, remaining) = parser.parse_known_args()
 
-  if remaining == ['help']:
-    parser.print_help()
-    parser.exit()
-
-  elif remaining == ['version']:
+def show_version(_, __, value):
+  if value:
     common_config.print_build_info()
-    parser.exit()
-
-  elif remaining != []:
-    Log.error('Invalid subcommand')
-    sys.exit(1)
-
-  namespace = vars(args)
-
-  log.set_logging_level(namespace)
+    sys.exit(0)
+
+
+@click.command()
+@click.option(
+    "--version",
+    is_flag=True,
+    is_eager=True,
+    expose_value=False,
+    callback=show_version,
+)
+@click.option('--verbose', is_flag=True)
+@click.option(
+    '--config-file',
+    help="path to a tracker config file",
+    default=os.path.join(utils.get_heron_tracker_conf_dir(), constants.DEFAULT_CONFIG_FILE),
+    show_default=True,
+)
+@click.option(
+    '--port',
+    type=int,
+    default=constants.DEFAULT_PORT,
+    show_default=True,
+    help="local port to serve on",
+)
+@click.option(
+    '--type',
+    "stmgr_type",
+    help=f"statemanager type e.g. {constants.DEFAULT_STATE_MANAGER_TYPE}",
+    type=click.Choice(choices=["file", "zookeeper"]),
+)
+@click.option(
+    '--name',
+    help=f"statemanager name e.g. {constants.DEFAULT_STATE_MANAGER_NAME}",
+)
+@click.option(
+    '--rootpath',
+    help=f"statemanager rootpath e.g. {constants.DEFAULT_STATE_MANAGER_ROOTPATH}",
+)
+@click.option(
+    '--tunnelhost',
+    help=f"statemanager tunnelhost e.g. {constants.DEFAULT_STATE_MANAGER_TUNNELHOST}",
+)
+@click.option(
+    '--hostport',
+    help=f"statemanager hostport e.g. {constants.DEFAULT_STATE_MANAGER_HOSTPORT}",
+)
+def cli(
+    config_file: str,
+    stmgr_type: str,
+    name: str,
+    rootpath: str,
+    tunnelhost: str,
+    hostport: str,
+    port: int,
+    verbose: bool,
+) -> None:
+  """
+  A HTTP service for serving data about clusters.
+
+  The statemanager's config from the given config file can be overrided using
+  options on this executable.
+
+  """
+
+  log.configure(logging.DEBUG if verbose else logging.INFO)
 
   # set Tornado global option
-  define_options(namespace['port'], namespace['config_file'])
+  define_options(port, config_file)
 
-  config = Config(create_tracker_config(namespace))
+  stmgr_override = {
+      "type": stmgr_type,
+      "name": name,
+      "rootpath": rootpath,
+      "tunnelhost": tunnelhost,
+      "hostport": hostport,
+  }
+  config = Config(create_tracker_config(config_file, stmgr_override))
 
   # create Tornado application
   application = Application(config)
@@ -278,15 +217,15 @@ def signal_handler(signum, frame):
   signal.signal(signal.SIGINT, signal_handler)
   signal.signal(signal.SIGTERM, signal_handler)
 
-  Log.info("Running on port: %d", namespace['port'])
-  if namespace["config_file"]:
-    Log.info("Using config file: %s", namespace['config_file'])
-  Log.info("Using state manager:\n" + str(config))
+  Log.info("Running on port: %d", port)
+  if config_file:
+    Log.info("Using config file: %s", config_file)
+  Log.info(f"Using state manager:\n{config}")
 
   http_server = tornado.httpserver.HTTPServer(application)
-  http_server.listen(namespace['port'])
+  http_server.listen(port)
 
   tornado.ioloop.IOLoop.instance().start()
 
 if __name__ == "__main__":
-  main()
+  cli() # pylint: disable=no-value-for-parameter

heron/tools/tracker/src/python/utils.py

@@ -162,7 +162,6 @@ def parse_config_file(config_file):
   if not os.path.lexists(expanded_config_file_path):
     return None
 
-  configs = {}
   # Read the configuration file
   with open(expanded_config_file_path, 'r') as f:
     configs = yaml.load(f)