catapult/third_party/gsutil/third_party/apitools/apitools/base/py/base_cli.py

#!/usr/bin/env python
"""Base script for generated CLI."""

import atexit
import code
import logging
import os
import readline
import rlcompleter
import sys

from google.apputils import appcommands
import gflags as flags

from apitools.base.py import encoding
from apitools.base.py import exceptions

__all__ = [
    'ConsoleWithReadline',
    'DeclareBaseFlags',
    'FormatOutput',
    'SetupLogger',
    'run_main',
]


# TODO(craigcitro): We should move all the flags for the
# StandardQueryParameters into this file, so that they can be used
# elsewhere easily.

_BASE_FLAGS_DECLARED = False
_OUTPUT_FORMATTER_MAP = {
    'protorpc': lambda x: x,
    'json': encoding.MessageToJson,
}


def DeclareBaseFlags():
    """Declare base flags for all CLIs."""
    # TODO(craigcitro): FlagValidators?
    global _BASE_FLAGS_DECLARED  # pylint: disable=global-statement
    if _BASE_FLAGS_DECLARED:
        return
    flags.DEFINE_boolean(
        'log_request', False,
        'Log requests.')
    flags.DEFINE_boolean(
        'log_response', False,
        'Log responses.')
    flags.DEFINE_boolean(
        'log_request_response', False,
        'Log requests and responses.')
    flags.DEFINE_enum(
        'output_format',
        'protorpc',
        _OUTPUT_FORMATTER_MAP.keys(),
        'Display format for results.')

    _BASE_FLAGS_DECLARED = True

# NOTE: This is specified here so that it can be read by other files
# without depending on the flag to be registered.
TRACE_HELP = (
    'A tracing token of the form "token:<tokenid>" '
    'to include in api requests.')
FLAGS = flags.FLAGS


def SetupLogger():
    if FLAGS.log_request or FLAGS.log_response or FLAGS.log_request_response:
        logging.basicConfig()
        logging.getLogger().setLevel(logging.INFO)


def FormatOutput(message, output_format=None):
    """Convert the output to the user-specified format."""
    output_format = output_format or FLAGS.output_format
    formatter = _OUTPUT_FORMATTER_MAP.get(FLAGS.output_format)
    if formatter is None:
        raise exceptions.UserError('Unknown output format: %s' % output_format)
    return formatter(message)


class _SmartCompleter(rlcompleter.Completer):

    def _callable_postfix(self, val, word):
        if ('(' in readline.get_line_buffer() or
                not callable(val)):
            return word
        else:
            return word + '('

    def complete(self, text, state):
        if not readline.get_line_buffer().strip():
            if not state:
                return '  '
            else:
                return None
        return rlcompleter.Completer.complete(self, text, state)


class ConsoleWithReadline(code.InteractiveConsole):

    """InteractiveConsole with readline, tab completion, and history."""

    def __init__(self, env, filename='<console>', histfile=None):
        new_locals = dict(env)
        new_locals.update({
            '_SmartCompleter': _SmartCompleter,
            'readline': readline,
            'rlcompleter': rlcompleter,
        })
        code.InteractiveConsole.__init__(self, new_locals, filename)
        readline.parse_and_bind('tab: complete')
        readline.set_completer(_SmartCompleter(new_locals).complete)
        if histfile is not None:
            histfile = os.path.expanduser(histfile)
            if os.path.exists(histfile):
                readline.read_history_file(histfile)
            atexit.register(lambda: readline.write_history_file(histfile))


def run_main():  # pylint: disable=invalid-name
    """Function to be used as setuptools script entry point.

    Appcommands assumes that it always runs as __main__, but launching
    via a setuptools-generated entry_point breaks this rule. We do some
    trickery here to make sure that appcommands and flags find their
    state where they expect to by faking ourselves as __main__.
    """

    # Put the flags for this module somewhere the flags module will look
    # for them.
    # pylint: disable=protected-access
    new_name = flags._GetMainModule()
    sys.modules[new_name] = sys.modules['__main__']
    for flag in FLAGS.FlagsByModuleDict().get(__name__, []):
        FLAGS._RegisterFlagByModule(new_name, flag)
        for key_flag in FLAGS.KeyFlagsByModuleDict().get(__name__, []):
            FLAGS._RegisterKeyFlagForModule(new_name, key_flag)
    # pylint: enable=protected-access

    # Now set __main__ appropriately so that appcommands will be
    # happy.
    sys.modules['__main__'] = sys.modules[__name__]
    appcommands.Run()
    sys.modules['__main__'] = sys.modules.pop(new_name)


if __name__ == '__main__':
    appcommands.Run()