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):
  17     self.host = host
  18     self.port = port
  19     self.handle = None
  20
  21   def setHandle(self, h):
  22     self.handle = h
  23
  24   def isOpen(self):
  25     return self.handle != None
  26
  27   def setTimeout(self, ms):
  28     if (self.handle != None):
  29       self.handle.settimeout(ms/1000.00)
  30     else:
  31       raise TTransportException(TTransportException.NOT_OPEN, 'No handle yet in TSocket')
  32
  33   def open(self):
  34     try:
  35       res0 = socket.getaddrinfo(self.host, self.port, socket.AF_UNSPEC, socket.SOCK_STREAM, 0, socket.AI_PASSIVE | socket.AI_ADDRCONFIG)
  36       for res in res0:
  37         self.handle = socket.socket(res[0], res[1])
  38         try:
  39           self.handle.connect(res[4])
  40         except socket.error, e:
  41           if res is not res0[-1]:
  42             continue
  43           else:
  44             raise e
  45         break
  46     except socket.error, e:
  47       raise TTransportException(TTransportException.NOT_OPEN, 'Could not connect to %s:%d' % (self.host, self.port))
  48
  49   def close(self):
  50     if self.handle != None:
  51       self.handle.close()
  52       self.handle = None
  53
  54   def read(self, sz):
  55     buff = self.handle.recv(sz)
  56     if len(buff) == 0:
  57       raise TTransportException('TSocket read 0 bytes')
  58     return buff
  59
  60   def write(self, buff):
  61     sent = 0
  62     have = len(buff)
  63     while sent < have:
  64       plus = self.handle.send(buff)
  65       if plus == 0:
  66         raise TTransportException('TSocket sent 0 bytes')
  67       sent += plus
  68       buff = buff[plus:]
  69
  70   def flush(self):
  71     pass
  72
  73 class TServerSocket(TServerTransportBase):
  74
  75   """Socket implementation of TServerTransport base."""
  76
  77   def __init__(self, port):
  78     self.port = port
  79     self.handle = None
  80
  81   def listen(self):
  82     res0 = socket.getaddrinfo(None, self.port, socket.AF_UNSPEC, socket.SOCK_STREAM, 0, socket.AI_PASSIVE | socket.AI_ADDRCONFIG)
  83     for res in res0:
  84       if res[0] is socket.AF_INET6 or res is res0[-1]:
  85         break
  86
  87     self.handle = socket.socket(res[0], res[1])
  88     self.handle.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  89     if hasattr(self.handle, 'set_timeout'):
  90       self.handle.set_timeout(None)
  91     self.handle.bind(res[4])
  92     self.handle.listen(128)
  93
  94   def accept(self):
  95     (client, addr) = self.handle.accept()
  96     result = TSocket()
  97     result.setHandle(client)
  98     return result
  99
 100   def close(self):
 101     self.handle.close()
 102     self.handle = None