lib/py/src/transport/TSocket.py

   1 #!/usr/bin/env python
   2 #
   3 # Copyright (c) 2006- Facebook
   4 # Distributed under the Thrift Software License
   5 #
   6 # See accompanying file LICENSE or visit the Thrift site at:
   7 # http://developers.facebook.com/thrift/
   8
   9 from TTransport import *
  10 import socket
  11
  12 class TSocket(TTransportBase):
  13
  14   """Socket implementation of TTransport base."""
  15
  16   def __init__(self, host='localhost', port=9090, unix_socket=None):
  17     """Initialize a TSocket
  18
  19     @param host(str)  The host to connect to.
  20     @param port(int)  The (TCP) port to connect to.
  21     @param unix_socket(str)  The filename of a unix socket to connect to.
  22                              (host and port will be ignored.)
  23     """
  24
  25     self.host = host
  26     self.port = port
  27     self.handle = None
  28     self._unix_socket = unix_socket
  29     self._timeout = None
  30
  31   def setHandle(self, h):
  32     self.handle = h
  33
  34   def isOpen(self):
  35     return self.handle != None
  36
  37   def setTimeout(self, ms):
  38     if ms is None:
  39       self._timeout = None
  40     else:
  41       self._timeout = ms/1000.0
  42
  43     if (self.handle != None):
  44       self.handle.settimeout(self._timeout)
  45
  46   def _resolveAddr(self):
  47     if self._unix_socket is not None:
  48       return [(socket.AF_UNIX, socket.SOCK_STREAM, None, None, self._unix_socket)]
  49     else:
  50       return socket.getaddrinfo(self.host, self.port, socket.AF_UNSPEC, socket.SOCK_STREAM, 0, socket.AI_PASSIVE | socket.AI_ADDRCONFIG)
  51
  52   def open(self):
  53     try:
  54       res0 = self._resolveAddr()
  55       for res in res0:
  56         self.handle = socket.socket(res[0], res[1])
  57         self.handle.settimeout(self._timeout)
  58         try:
  59           self.handle.connect(res[4])
  60         except socket.error, e:
  61           if res is not res0[-1]:
  62             continue
  63           else:
  64             raise e
  65         break
  66     except socket.error, e:
  67       raise TTransportException(TTransportException.NOT_OPEN, 'Could not connect to %s:%d' % (self.host, self.port))
  68
  69   def close(self):
  70     if self.handle != None:
  71       self.handle.close()
  72       self.handle = None
  73
  74   def read(self, sz):
  75     buff = self.handle.recv(sz)
  76     if len(buff) == 0:
  77       raise TTransportException('TSocket read 0 bytes')
  78     return buff
  79
  80   def write(self, buff):
  81     sent = 0
  82     have = len(buff)
  83     while sent < have:
  84       plus = self.handle.send(buff)
  85       if plus == 0:
  86         raise TTransportException('TSocket sent 0 bytes')
  87       sent += plus
  88       buff = buff[plus:]
  89
  90   def flush(self):
  91     pass
  92
  93 class TServerSocket(TServerTransportBase):
  94
  95   """Socket implementation of TServerTransport base."""
  96
  97   def __init__(self, port):
  98     self.port = port
  99     self.handle = None
 100
 101   def listen(self):
 102     res0 = socket.getaddrinfo(None, self.port, socket.AF_UNSPEC, socket.SOCK_STREAM, 0, socket.AI_PASSIVE | socket.AI_ADDRCONFIG)
 103     for res in res0:
 104       if res[0] is socket.AF_INET6 or res is res0[-1]:
 105         break
 106
 107     self.handle = socket.socket(res[0], res[1])
 108     self.handle.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 109     if hasattr(self.handle, 'set_timeout'):
 110       self.handle.set_timeout(None)
 111     self.handle.bind(res[4])
 112     self.handle.listen(128)
 113
 114   def accept(self):
 115     (client, addr) = self.handle.accept()
 116     result = TSocket()
 117     result.setHandle(client)
 118     return result
 119
 120   def close(self):
 121     self.handle.close()
 122     self.handle = None