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