1#!/usr/bin/env python
2"""Base script for generated CLI."""
3
4import atexit
5import code
6import logging
7import os
8import readline
9import rlcompleter
10import sys
11
12from google.apputils import appcommands
13import gflags as flags
14
15from apitools.base.py import encoding
16from apitools.base.py import exceptions
17
18__all__ = [
19    'ConsoleWithReadline',
20    'DeclareBaseFlags',
21    'FormatOutput',
22    'SetupLogger',
23    'run_main',
24]
25
26
27# TODO(craigcitro): We should move all the flags for the
28# StandardQueryParameters into this file, so that they can be used
29# elsewhere easily.
30
31_BASE_FLAGS_DECLARED = False
32_OUTPUT_FORMATTER_MAP = {
33    'protorpc': lambda x: x,
34    'json': encoding.MessageToJson,
35}
36
37
38def DeclareBaseFlags():
39    """Declare base flags for all CLIs."""
40    # TODO(craigcitro): FlagValidators?
41    global _BASE_FLAGS_DECLARED  # pylint: disable=global-statement
42    if _BASE_FLAGS_DECLARED:
43        return
44    flags.DEFINE_boolean(
45        'log_request', False,
46        'Log requests.')
47    flags.DEFINE_boolean(
48        'log_response', False,
49        'Log responses.')
50    flags.DEFINE_boolean(
51        'log_request_response', False,
52        'Log requests and responses.')
53    flags.DEFINE_enum(
54        'output_format',
55        'protorpc',
56        _OUTPUT_FORMATTER_MAP.keys(),
57        'Display format for results.')
58
59    _BASE_FLAGS_DECLARED = True
60
61# NOTE: This is specified here so that it can be read by other files
62# without depending on the flag to be registered.
63TRACE_HELP = (
64    'A tracing token of the form "token:<tokenid>" '
65    'to include in api requests.')
66FLAGS = flags.FLAGS
67
68
69def SetupLogger():
70    if FLAGS.log_request or FLAGS.log_response or FLAGS.log_request_response:
71        logging.basicConfig()
72        logging.getLogger().setLevel(logging.INFO)
73
74
75def FormatOutput(message, output_format=None):
76    """Convert the output to the user-specified format."""
77    output_format = output_format or FLAGS.output_format
78    formatter = _OUTPUT_FORMATTER_MAP.get(FLAGS.output_format)
79    if formatter is None:
80        raise exceptions.UserError('Unknown output format: %s' % output_format)
81    return formatter(message)
82
83
84class _SmartCompleter(rlcompleter.Completer):
85
86    def _callable_postfix(self, val, word):
87        if ('(' in readline.get_line_buffer() or
88                not callable(val)):
89            return word
90        else:
91            return word + '('
92
93    def complete(self, text, state):
94        if not readline.get_line_buffer().strip():
95            if not state:
96                return '  '
97            else:
98                return None
99        return rlcompleter.Completer.complete(self, text, state)
100
101
102class ConsoleWithReadline(code.InteractiveConsole):
103
104    """InteractiveConsole with readline, tab completion, and history."""
105
106    def __init__(self, env, filename='<console>', histfile=None):
107        new_locals = dict(env)
108        new_locals.update({
109            '_SmartCompleter': _SmartCompleter,
110            'readline': readline,
111            'rlcompleter': rlcompleter,
112        })
113        code.InteractiveConsole.__init__(self, new_locals, filename)
114        readline.parse_and_bind('tab: complete')
115        readline.set_completer(_SmartCompleter(new_locals).complete)
116        if histfile is not None:
117            histfile = os.path.expanduser(histfile)
118            if os.path.exists(histfile):
119                readline.read_history_file(histfile)
120            atexit.register(lambda: readline.write_history_file(histfile))
121
122
123def run_main():  # pylint: disable=invalid-name
124    """Function to be used as setuptools script entry point.
125
126    Appcommands assumes that it always runs as __main__, but launching
127    via a setuptools-generated entry_point breaks this rule. We do some
128    trickery here to make sure that appcommands and flags find their
129    state where they expect to by faking ourselves as __main__.
130    """
131
132    # Put the flags for this module somewhere the flags module will look
133    # for them.
134    # pylint: disable=protected-access
135    new_name = flags._GetMainModule()
136    sys.modules[new_name] = sys.modules['__main__']
137    for flag in FLAGS.FlagsByModuleDict().get(__name__, []):
138        FLAGS._RegisterFlagByModule(new_name, flag)
139        for key_flag in FLAGS.KeyFlagsByModuleDict().get(__name__, []):
140            FLAGS._RegisterKeyFlagForModule(new_name, key_flag)
141    # pylint: enable=protected-access
142
143    # Now set __main__ appropriately so that appcommands will be
144    # happy.
145    sys.modules['__main__'] = sys.modules[__name__]
146    appcommands.Run()
147    sys.modules['__main__'] = sys.modules.pop(new_name)
148
149
150if __name__ == '__main__':
151    appcommands.Run()
152