lib/py/src/transport/TTransport.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 cStringIO import StringIO
  10 from struct import pack,unpack
  11 from thrift.Thrift import TException
  12
  13 class TTransportException(TException):
  14
  15   """Custom Transport Exception class"""
  16
  17   UNKNOWN = 0
  18   NOT_OPEN = 1
  19   ALREADY_OPEN = 2
  20   TIMED_OUT = 3
  21   END_OF_FILE = 4
  22
  23   def __init__(self, type=UNKNOWN, message=None):
  24     TException.__init__(self, message)
  25     self.type = type
  26
  27 class TTransportBase:
  28
  29   """Base class for Thrift transport layer."""
  30
  31   def isOpen(self):
  32     pass
  33
  34   def open(self):
  35     pass
  36
  37   def close(self):
  38     pass
  39
  40   def read(self, sz):
  41     pass
  42
  43   def readAll(self, sz):
  44     buff = ''
  45     have = 0
  46     while (have < sz):
  47       chunk = self.read(sz-have)
  48       have += len(chunk)
  49       buff += chunk
  50
  51       if len(chunk) == 0:
  52         raise EOFError()
  53
  54     return buff
  55
  56   def write(self, buf):
  57     pass
  58
  59   def flush(self):
  60     pass
  61
  62 # This class should be thought of as an interface.
  63 class CReadableTransport:
  64   """base class for transports that are readable from C"""
  65
  66   # TODO(dreiss): Think about changing this interface to allow us to use
  67   #               a (Python, not c) StringIO instead, because it allows
  68   #               you to write after reading.
  69
  70   # NOTE: This is a classic class, so properties will NOT work
  71   #       correctly for setting.
  72   @property
  73   def cstringio_buf(self):
  74     """A cStringIO buffer that contains the current chunk we are reading."""
  75     pass
  76
  77   def cstringio_refill(self, partialread, reqlen):
  78     """Refills cstringio_buf.
  79
  80     Returns the currently used buffer (which can but need not be the same as
  81     the old cstringio_buf). partialread is what the C code has read from the
  82     buffer, and should be inserted into the buffer before any more reads.  The
  83     return value must be a new, not borrowed reference.  Something along the
  84     lines of self._buf should be fine.
  85
  86     If reqlen bytes can't be read, throw EOFError.
  87     """
  88     pass
  89
  90 class TServerTransportBase:
  91
  92   """Base class for Thrift server transports."""
  93
  94   def listen(self):
  95     pass
  96
  97   def accept(self):
  98     pass
  99
 100   def close(self):
 101     pass
 102
 103 class TTransportFactoryBase:
 104
 105   """Base class for a Transport Factory"""
 106
 107   def getTransport(self, trans):
 108     return trans
 109
 110 class TBufferedTransportFactory:
 111
 112   """Factory transport that builds buffered transports"""
 113
 114   def getTransport(self, trans):
 115     buffered = TBufferedTransport(trans)
 116     return buffered
 117
 118
 119 class TBufferedTransport(TTransportBase,CReadableTransport):
 120
 121   """Class that wraps another transport and buffers its I/O."""
 122
 123   DEFAULT_BUFFER = 4096
 124
 125   def __init__(self, trans):
 126     self.__trans = trans
 127     self.__wbuf = StringIO()
 128     self.__rbuf = StringIO("")
 129
 130   def isOpen(self):
 131     return self.__trans.isOpen()
 132
 133   def open(self):
 134     return self.__trans.open()
 135
 136   def close(self):
 137     return self.__trans.close()
 138
 139   def read(self, sz):
 140     ret = self.__rbuf.read(sz)
 141     if len(ret) != 0:
 142       return ret
 143
 144     self.__rbuf = StringIO(self.__trans.read(max(sz, self.DEFAULT_BUFFER)))
 145     return self.__rbuf.read(sz)
 146
 147   def write(self, buf):
 148     self.__wbuf.write(buf)
 149
 150   def flush(self):
 151     self.__trans.write(self.__wbuf.getvalue())
 152     self.__trans.flush()
 153     self.__wbuf = StringIO()
 154
 155   # Implement the CReadableTransport interface.
 156   @property
 157   def cstringio_buf(self):
 158     return self.__rbuf
 159
 160   def cstringio_refill(self, partialread, reqlen):
 161     retstring = partialread
 162     if reqlen < self.DEFAULT_BUFFER:
 163       # try to make a read of as much as we can.
 164       retstring += self.__trans.read(self.DEFAULT_BUFFER)
 165
 166     # but make sure we do read reqlen bytes.
 167     if len(retstring) < reqlen:
 168       retstring += self.__trans.readAll(reqlen - len(retstring))
 169
 170     self.__rbuf = StringIO(retstring)
 171     return self.__rbuf
 172
 173 class TMemoryBuffer(TTransportBase, CReadableTransport):
 174   """Wraps a cStringIO object as a TTransport.
 175
 176   NOTE: Unlike the C++ version of this class, you cannot write to it
 177         then immediately read from it.  If you want to read from a
 178         TMemoryBuffer, you must either pass a string to the constructor.
 179   TODO(dreiss): Make this work like the C++ version.
 180   """
 181
 182   def __init__(self, value=None):
 183     """value -- a value to read from for stringio
 184
 185     If value is set, this will be a transport for reading,
 186     otherwise, it is for writing"""
 187     if value is not None:
 188       self._buffer = StringIO(value)
 189     else:
 190       self._buffer = StringIO()
 191
 192   def isOpen(self):
 193     return not self._buffer.closed
 194
 195   def open(self):
 196     pass
 197
 198   def close(self):
 199     self._buffer.close()
 200
 201   def read(self, sz):
 202     return self._buffer.read(sz)
 203
 204   def write(self, buf):
 205     self._buffer.write(buf)
 206
 207   def flush(self):
 208     pass
 209
 210   def getvalue(self):
 211     return self._buffer.getvalue()
 212
 213   # Implement the CReadableTransport interface.
 214   @property
 215   def cstringio_buf(self):
 216     return self._buffer
 217
 218   def cstringio_refill(self, partialread, reqlen):
 219     # only one shot at reading...
 220     raise EOFError()
 221
 222 class TFramedTransportFactory:
 223
 224   """Factory transport that builds framed transports"""
 225
 226   def getTransport(self, trans):
 227     framed = TFramedTransport(trans)
 228     return framed
 229
 230
 231 class TFramedTransport(TTransportBase):
 232
 233   """Class that wraps another transport and frames its I/O when writing."""
 234
 235   def __init__(self, trans, read=True, write=True):
 236     self.__trans = trans
 237     if read:
 238       self.__rbuf = ''
 239     else:
 240       self.__rbuf = None
 241     if write:
 242       self.__wbuf = StringIO()
 243     else:
 244       self.__wbuf = None
 245
 246   def isOpen(self):
 247     return self.__trans.isOpen()
 248
 249   def open(self):
 250     return self.__trans.open()
 251
 252   def close(self):
 253     return self.__trans.close()
 254
 255   def read(self, sz):
 256     if self.__rbuf == None:
 257       return self.__trans.read(sz)
 258     if len(self.__rbuf) == 0:
 259       self.readFrame()
 260     give = min(len(self.__rbuf), sz)
 261     buff = self.__rbuf[0:give]
 262     self.__rbuf = self.__rbuf[give:]
 263     return buff
 264
 265   def readFrame(self):
 266     buff = self.__trans.readAll(4)
 267     sz, = unpack('!i', buff)
 268     self.__rbuf = self.__trans.readAll(sz)
 269
 270   def write(self, buf):
 271     if self.__wbuf == None:
 272       return self.__trans.write(buf)
 273     self.__wbuf.write(buf)
 274
 275   def flush(self):
 276     if self.__wbuf == None:
 277       return self.__trans.flush()
 278     wout = self.__wbuf.getvalue()
 279     wsz = len(wout)
 280     # N.B.: Doing this string concatenation is WAY cheaper than making
 281     # two separate calls to the underlying socket object. Socket writes in
 282     # Python turn out to be REALLY expensive, but it seems to do a pretty
 283     # good job of managing string buffer operations without excessive copies
 284     buf = pack("!i", wsz) + wout
 285     self.__trans.write(buf)
 286     self.__trans.flush()
 287     self.__wbuf = StringIO()