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