1import argparse
2import socket
3
4from local_thrift import thrift  # noqa
5from thrift.transport.TSocket import TSocket
6from thrift.transport.TTransport import TBufferedTransport, TFramedTransport
7from thrift.transport.THttpClient import THttpClient
8from thrift.protocol.TBinaryProtocol import TBinaryProtocol
9from thrift.protocol.TCompactProtocol import TCompactProtocol
10from thrift.protocol.TJSONProtocol import TJSONProtocol
11
12
13def add_common_args(p):
14    p.add_argument('--host', default='localhost')
15    p.add_argument('--port', type=int, default=9090)
16    p.add_argument('--protocol', default='binary')
17    p.add_argument('--transport', default='buffered')
18    p.add_argument('--ssl', action='store_true')
19
20
21def parse_common_args(argv):
22    p = argparse.ArgumentParser()
23    add_common_args(p)
24    return p.parse_args(argv)
25
26
27def init_protocol(args):
28    sock = TSocket(args.host, args.port, socket_family=socket.AF_INET)
29    sock.setTimeout(500)
30    trans = {
31        'buffered': TBufferedTransport,
32        'framed': TFramedTransport,
33        'http': THttpClient,
34    }[args.transport](sock)
35    trans.open()
36    return {
37        'binary': TBinaryProtocol,
38        'compact': TCompactProtocol,
39        'json': TJSONProtocol,
40    }[args.protocol](trans)
41