]>
Commit | Line | Data |
---|---|---|
f67539c2 TL |
1 | import argparse |
2 | import socket | |
3 | ||
4 | from local_thrift import thrift # noqa | |
5 | from thrift.transport.TSocket import TSocket | |
6 | from thrift.transport.TTransport import TBufferedTransport, TFramedTransport | |
7 | from thrift.transport.THttpClient import THttpClient | |
8 | from thrift.protocol.TBinaryProtocol import TBinaryProtocol | |
9 | from thrift.protocol.TCompactProtocol import TCompactProtocol | |
10 | from thrift.protocol.TJSONProtocol import TJSONProtocol | |
11 | ||
12 | ||
13 | def 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 | ||
21 | def parse_common_args(argv): | |
22 | p = argparse.ArgumentParser() | |
23 | add_common_args(p) | |
24 | return p.parse_args(argv) | |
25 | ||
26 | ||
27 | def 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) |