lib/thrift/protocol/TProtocol.py

   1 #
   2 # Licensed to the Apache Software Foundation (ASF) under one
   3 # or more contributor license agreements. See the NOTICE file
   4 # distributed with this work for additional information
   5 # regarding copyright ownership. The ASF licenses this file
   6 # to you under the Apache License, Version 2.0 (the
   7 # "License"); you may not use this file except in compliance
   8 # with the License. You may obtain a copy of the License at
   9 #
  10 #   http://www.apache.org/licenses/LICENSE-2.0
  11 #
  12 # Unless required by applicable law or agreed to in writing,
  13 # software distributed under the License is distributed on an
  14 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  15 # KIND, either express or implied. See the License for the
  16 # specific language governing permissions and limitations
  17 # under the License.
  18 #
  19
  20 from thrift.Thrift import *
  21
  22
  23 class TProtocolException(TException):
  24   """Custom Protocol Exception class"""
  25
  26   UNKNOWN = 0
  27   INVALID_DATA = 1
  28   NEGATIVE_SIZE = 2
  29   SIZE_LIMIT = 3
  30   BAD_VERSION = 4
  31   NOT_IMPLEMENTED = 5
  32   DEPTH_LIMIT = 6
  33
  34   def __init__(self, type=UNKNOWN, message=None):
  35     TException.__init__(self, message)
  36     self.type = type
  37
  38
  39 class TProtocolBase:
  40   """Base class for Thrift protocol driver."""
  41
  42   def __init__(self, trans):
  43     self.trans = trans
  44
  45   def writeMessageBegin(self, name, ttype, seqid):
  46     pass
  47
  48   def writeMessageEnd(self):
  49     pass
  50
  51   def writeStructBegin(self, name):
  52     pass
  53
  54   def writeStructEnd(self):
  55     pass
  56
  57   def writeFieldBegin(self, name, ttype, fid):
  58     pass
  59
  60   def writeFieldEnd(self):
  61     pass
  62
  63   def writeFieldStop(self):
  64     pass
  65
  66   def writeMapBegin(self, ktype, vtype, size):
  67     pass
  68
  69   def writeMapEnd(self):
  70     pass
  71
  72   def writeListBegin(self, etype, size):
  73     pass
  74
  75   def writeListEnd(self):
  76     pass
  77
  78   def writeSetBegin(self, etype, size):
  79     pass
  80
  81   def writeSetEnd(self):
  82     pass
  83
  84   def writeBool(self, bool_val):
  85     pass
  86
  87   def writeByte(self, byte):
  88     pass
  89
  90   def writeI16(self, i16):
  91     pass
  92
  93   def writeI32(self, i32):
  94     pass
  95
  96   def writeI64(self, i64):
  97     pass
  98
  99   def writeDouble(self, dub):
 100     pass
 101
 102   def writeString(self, str_val):
 103     pass
 104
 105   def readMessageBegin(self):
 106     pass
 107
 108   def readMessageEnd(self):
 109     pass
 110
 111   def readStructBegin(self):
 112     pass
 113
 114   def readStructEnd(self):
 115     pass
 116
 117   def readFieldBegin(self):
 118     pass
 119
 120   def readFieldEnd(self):
 121     pass
 122
 123   def readMapBegin(self):
 124     pass
 125
 126   def readMapEnd(self):
 127     pass
 128
 129   def readListBegin(self):
 130     pass
 131
 132   def readListEnd(self):
 133     pass
 134
 135   def readSetBegin(self):
 136     pass
 137
 138   def readSetEnd(self):
 139     pass
 140
 141   def readBool(self):
 142     pass
 143
 144   def readByte(self):
 145     pass
 146
 147   def readI16(self):
 148     pass
 149
 150   def readI32(self):
 151     pass
 152
 153   def readI64(self):
 154     pass
 155
 156   def readDouble(self):
 157     pass
 158
 159   def readString(self):
 160     pass
 161
 162   def skip(self, ttype):
 163     if ttype == TType.STOP:
 164       return
 165     elif ttype == TType.BOOL:
 166       self.readBool()
 167     elif ttype == TType.BYTE:
 168       self.readByte()
 169     elif ttype == TType.I16:
 170       self.readI16()
 171     elif ttype == TType.I32:
 172       self.readI32()
 173     elif ttype == TType.I64:
 174       self.readI64()
 175     elif ttype == TType.DOUBLE:
 176       self.readDouble()
 177     elif ttype == TType.STRING:
 178       self.readString()
 179     elif ttype == TType.STRUCT:
 180       name = self.readStructBegin()
 181       while True:
 182         (name, ttype, id) = self.readFieldBegin()
 183         if ttype == TType.STOP:
 184           break
 185         self.skip(ttype)
 186         self.readFieldEnd()
 187       self.readStructEnd()
 188     elif ttype == TType.MAP:
 189       (ktype, vtype, size) = self.readMapBegin()
 190       for i in xrange(size):
 191         self.skip(ktype)
 192         self.skip(vtype)
 193       self.readMapEnd()
 194     elif ttype == TType.SET:
 195       (etype, size) = self.readSetBegin()
 196       for i in xrange(size):
 197         self.skip(etype)
 198       self.readSetEnd()
 199     elif ttype == TType.LIST:
 200       (etype, size) = self.readListBegin()
 201       for i in xrange(size):
 202         self.skip(etype)
 203       self.readListEnd()
 204
 205   # tuple of: ( 'reader method' name, is_container bool, 'writer_method' name )
 206   _TTYPE_HANDLERS = (
 207        (None, None, False),  # 0 TType.STOP
 208        (None, None, False),  # 1 TType.VOID # TODO: handle void?
 209        ('readBool', 'writeBool', False),  # 2 TType.BOOL
 210        ('readByte',  'writeByte', False),  # 3 TType.BYTE and I08
 211        ('readDouble', 'writeDouble', False),  # 4 TType.DOUBLE
 212        (None, None, False),  # 5 undefined
 213        ('readI16', 'writeI16', False),  # 6 TType.I16
 214        (None, None, False),  # 7 undefined
 215        ('readI32', 'writeI32', False),  # 8 TType.I32
 216        (None, None, False),  # 9 undefined
 217        ('readI64', 'writeI64', False),  # 10 TType.I64
 218        ('readString', 'writeString', False),  # 11 TType.STRING and UTF7
 219        ('readContainerStruct', 'writeContainerStruct', True),  # 12 *.STRUCT
 220        ('readContainerMap', 'writeContainerMap', True),  # 13 TType.MAP
 221        ('readContainerSet', 'writeContainerSet', True),  # 14 TType.SET
 222        ('readContainerList', 'writeContainerList', True),  # 15 TType.LIST
 223        (None, None, False),  # 16 TType.UTF8 # TODO: handle utf8 types?
 224        (None, None, False)  # 17 TType.UTF16 # TODO: handle utf16 types?
 225       )
 226
 227   def readFieldByTType(self, ttype, spec):
 228     try:
 229       (r_handler, w_handler, is_container) = self._TTYPE_HANDLERS[ttype]
 230     except IndexError:
 231       raise TProtocolException(type=TProtocolException.INVALID_DATA,
 232                                message='Invalid field type %d' % (ttype))
 233     if r_handler is None:
 234       raise TProtocolException(type=TProtocolException.INVALID_DATA,
 235                                message='Invalid field type %d' % (ttype))
 236     reader = getattr(self, r_handler)
 237     if not is_container:
 238       return reader()
 239     return reader(spec)
 240
 241   def readContainerList(self, spec):
 242     results = []
 243     ttype, tspec = spec[0], spec[1]
 244     r_handler = self._TTYPE_HANDLERS[ttype][0]
 245     reader = getattr(self, r_handler)
 246     (list_type, list_len) = self.readListBegin()
 247     if tspec is None:
 248       # list values are simple types
 249       for idx in xrange(list_len):
 250         results.append(reader())
 251     else:
 252       # this is like an inlined readFieldByTType
 253       container_reader = self._TTYPE_HANDLERS[list_type][0]
 254       val_reader = getattr(self, container_reader)
 255       for idx in xrange(list_len):
 256         val = val_reader(tspec)
 257         results.append(val)
 258     self.readListEnd()
 259     return results
 260
 261   def readContainerSet(self, spec):
 262     results = set()
 263     ttype, tspec = spec[0], spec[1]
 264     r_handler = self._TTYPE_HANDLERS[ttype][0]
 265     reader = getattr(self, r_handler)
 266     (set_type, set_len) = self.readSetBegin()
 267     if tspec is None:
 268       # set members are simple types
 269       for idx in xrange(set_len):
 270         results.add(reader())
 271     else:
 272       container_reader = self._TTYPE_HANDLERS[set_type][0]
 273       val_reader = getattr(self, container_reader)
 274       for idx in xrange(set_len):
 275         results.add(val_reader(tspec))
 276     self.readSetEnd()
 277     return results
 278
 279   def readContainerStruct(self, spec):
 280     (obj_class, obj_spec) = spec
 281     obj = obj_class()
 282     obj.read(self)
 283     return obj
 284
 285   def readContainerMap(self, spec):
 286     results = dict()
 287     key_ttype, key_spec = spec[0], spec[1]
 288     val_ttype, val_spec = spec[2], spec[3]
 289     (map_ktype, map_vtype, map_len) = self.readMapBegin()
 290     # TODO: compare types we just decoded with thrift_spec and
 291     # abort/skip if types disagree
 292     key_reader = getattr(self, self._TTYPE_HANDLERS[key_ttype][0])
 293     val_reader = getattr(self, self._TTYPE_HANDLERS[val_ttype][0])
 294     # list values are simple types
 295     for idx in xrange(map_len):
 296       if key_spec is None:
 297         k_val = key_reader()
 298       else:
 299         k_val = self.readFieldByTType(key_ttype, key_spec)
 300       if val_spec is None:
 301         v_val = val_reader()
 302       else:
 303         v_val = self.readFieldByTType(val_ttype, val_spec)
 304       # this raises a TypeError with unhashable keys types
 305       # i.e. this fails: d=dict(); d[[0,1]] = 2
 306       results[k_val] = v_val
 307     self.readMapEnd()
 308     return results
 309
 310   def readStruct(self, obj, thrift_spec):
 311     self.readStructBegin()
 312     while True:
 313       (fname, ftype, fid) = self.readFieldBegin()
 314       if ftype == TType.STOP:
 315         break
 316       try:
 317         field = thrift_spec[fid]
 318       except IndexError:
 319         self.skip(ftype)
 320       else:
 321         if field is not None and ftype == field[1]:
 322           fname = field[2]
 323           fspec = field[3]
 324           val = self.readFieldByTType(ftype, fspec)
 325           setattr(obj, fname, val)
 326         else:
 327           self.skip(ftype)
 328       self.readFieldEnd()
 329     self.readStructEnd()
 330
 331   def writeContainerStruct(self, val, spec):
 332     val.write(self)
 333
 334   def writeContainerList(self, val, spec):
 335     self.writeListBegin(spec[0], len(val))
 336     r_handler, w_handler, is_container = self._TTYPE_HANDLERS[spec[0]]
 337     e_writer = getattr(self, w_handler)
 338     if not is_container:
 339       for elem in val:
 340         e_writer(elem)
 341     else:
 342       for elem in val:
 343         e_writer(elem, spec[1])
 344     self.writeListEnd()
 345
 346   def writeContainerSet(self, val, spec):
 347     self.writeSetBegin(spec[0], len(val))
 348     r_handler, w_handler, is_container = self._TTYPE_HANDLERS[spec[0]]
 349     e_writer = getattr(self, w_handler)
 350     if not is_container:
 351       for elem in val:
 352         e_writer(elem)
 353     else:
 354       for elem in val:
 355         e_writer(elem, spec[1])
 356     self.writeSetEnd()
 357
 358   def writeContainerMap(self, val, spec):
 359     k_type = spec[0]
 360     v_type = spec[2]
 361     ignore, ktype_name, k_is_container = self._TTYPE_HANDLERS[k_type]
 362     ignore, vtype_name, v_is_container = self._TTYPE_HANDLERS[v_type]
 363     k_writer = getattr(self, ktype_name)
 364     v_writer = getattr(self, vtype_name)
 365     self.writeMapBegin(k_type, v_type, len(val))
 366     for m_key, m_val in val.iteritems():
 367       if not k_is_container:
 368         k_writer(m_key)
 369       else:
 370         k_writer(m_key, spec[1])
 371       if not v_is_container:
 372         v_writer(m_val)
 373       else:
 374         v_writer(m_val, spec[3])
 375     self.writeMapEnd()
 376
 377   def writeStruct(self, obj, thrift_spec):
 378     self.writeStructBegin(obj.__class__.__name__)
 379     for field in thrift_spec:
 380       if field is None:
 381         continue
 382       fname = field[2]
 383       val = getattr(obj, fname)
 384       if val is None:
 385         # skip writing out unset fields
 386         continue
 387       fid = field[0]
 388       ftype = field[1]
 389       fspec = field[3]
 390       # get the writer method for this value
 391       self.writeFieldBegin(fname, ftype, fid)
 392       self.writeFieldByTType(ftype, val, fspec)
 393       self.writeFieldEnd()
 394     self.writeFieldStop()
 395     self.writeStructEnd()
 396
 397   def writeFieldByTType(self, ttype, val, spec):
 398     r_handler, w_handler, is_container = self._TTYPE_HANDLERS[ttype]
 399     writer = getattr(self, w_handler)
 400     if is_container:
 401       writer(val, spec)
 402     else:
 403       writer(val)
 404
 405 def checkIntegerLimits(i, bits):
 406     if bits == 8 and (i < -128 or i > 127):
 407         raise TProtocolException(TProtocolException.INVALID_DATA,
 408                                  "i8 requires -128 <= number <= 127")
 409     elif bits == 16 and (i < -32768 or i > 32767):
 410         raise TProtocolException(TProtocolException.INVALID_DATA,
 411                                  "i16 requires -32768 <= number <= 32767")
 412     elif bits == 32 and (i < -2147483648 or i > 2147483647):
 413         raise TProtocolException(TProtocolException.INVALID_DATA,
 414                                  "i32 requires -2147483648 <= number <= 2147483647")
 415     elif bits == 64 and (i < -9223372036854775808 or i > 9223372036854775807):
 416          raise TProtocolException(TProtocolException.INVALID_DATA,
 417                                   "i64 requires -9223372036854775808 <= number <= 9223372036854775807")
 418
 419 class TProtocolFactory:
 420   def getProtocol(self, trans):
 421     pass