lib/py/src/protocol/TProtocol.py

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#

from thrift.Thrift import TException, TType, TFrozenDict
from ..compat import binary_to_str, str_to_binary

import six
import sys
from itertools import islice
from six.moves import zip


class TProtocolException(TException):
  """Custom Protocol Exception class"""

  UNKNOWN = 0
  INVALID_DATA = 1
  NEGATIVE_SIZE = 2
  SIZE_LIMIT = 3
  BAD_VERSION = 4
  NOT_IMPLEMENTED = 5
  DEPTH_LIMIT = 6

  def __init__(self, type=UNKNOWN, message=None):
    TException.__init__(self, message)
    self.type = type


class TProtocolBase(object):
  """Base class for Thrift protocol driver."""

  def __init__(self, trans):
    self.trans = trans

  def writeMessageBegin(self, name, ttype, seqid):
    pass

  def writeMessageEnd(self):
    pass

  def writeStructBegin(self, name):
    pass

  def writeStructEnd(self):
    pass

  def writeFieldBegin(self, name, ttype, fid):
    pass

  def writeFieldEnd(self):
    pass

  def writeFieldStop(self):
    pass

  def writeMapBegin(self, ktype, vtype, size):
    pass

  def writeMapEnd(self):
    pass

  def writeListBegin(self, etype, size):
    pass

  def writeListEnd(self):
    pass

  def writeSetBegin(self, etype, size):
    pass

  def writeSetEnd(self):
    pass

  def writeBool(self, bool_val):
    pass

  def writeByte(self, byte):
    pass

  def writeI16(self, i16):
    pass

  def writeI32(self, i32):
    pass

  def writeI64(self, i64):
    pass

  def writeDouble(self, dub):
    pass

  def writeString(self, str_val):
    self.writeBinary(str_to_binary(str_val))

  def writeBinary(self, str_val):
    pass

  def writeUtf8(self, str_val):
    self.writeString(str_val.encode('utf8'))

  def readMessageBegin(self):
    pass

  def readMessageEnd(self):
    pass

  def readStructBegin(self):
    pass

  def readStructEnd(self):
    pass

  def readFieldBegin(self):
    pass

  def readFieldEnd(self):
    pass

  def readMapBegin(self):
    pass

  def readMapEnd(self):
    pass

  def readListBegin(self):
    pass

  def readListEnd(self):
    pass

  def readSetBegin(self):
    pass

  def readSetEnd(self):
    pass

  def readBool(self):
    pass

  def readByte(self):
    pass

  def readI16(self):
    pass

  def readI32(self):
    pass

  def readI64(self):
    pass

  def readDouble(self):
    pass

  def readString(self):
    return binary_to_str(self.readBinary())

  def readBinary(self):
    pass

  def readUtf8(self):
    return self.readString().decode('utf8')

  def skip(self, ttype):
    if ttype == TType.STOP:
      return
    elif ttype == TType.BOOL:
      self.readBool()
    elif ttype == TType.BYTE:
      self.readByte()
    elif ttype == TType.I16:
      self.readI16()
    elif ttype == TType.I32:
      self.readI32()
    elif ttype == TType.I64:
      self.readI64()
    elif ttype == TType.DOUBLE:
      self.readDouble()
    elif ttype == TType.STRING:
      self.readString()
    elif ttype == TType.STRUCT:
      name = self.readStructBegin()
      while True:
        (name, ttype, id) = self.readFieldBegin()
        if ttype == TType.STOP:
          break
        self.skip(ttype)
        self.readFieldEnd()
      self.readStructEnd()
    elif ttype == TType.MAP:
      (ktype, vtype, size) = self.readMapBegin()
      for i in range(size):
        self.skip(ktype)
        self.skip(vtype)
      self.readMapEnd()
    elif ttype == TType.SET:
      (etype, size) = self.readSetBegin()
      for i in range(size):
        self.skip(etype)
      self.readSetEnd()
    elif ttype == TType.LIST:
      (etype, size) = self.readListBegin()
      for i in range(size):
        self.skip(etype)
      self.readListEnd()

  # tuple of: ( 'reader method' name, is_container bool, 'writer_method' name )
  _TTYPE_HANDLERS = (
       (None, None, False),  # 0 TType.STOP
       (None, None, False),  # 1 TType.VOID # TODO: handle void?
       ('readBool', 'writeBool', False),  # 2 TType.BOOL
       ('readByte', 'writeByte', False),  # 3 TType.BYTE and I08
       ('readDouble', 'writeDouble', False),  # 4 TType.DOUBLE
       (None, None, False),  # 5 undefined
       ('readI16', 'writeI16', False),  # 6 TType.I16
       (None, None, False),  # 7 undefined
       ('readI32', 'writeI32', False),  # 8 TType.I32
       (None, None, False),  # 9 undefined
       ('readI64', 'writeI64', False),  # 10 TType.I64
       ('readString', 'writeString', False),  # 11 TType.STRING and UTF7
       ('readContainerStruct', 'writeContainerStruct', True),  # 12 *.STRUCT
       ('readContainerMap', 'writeContainerMap', True),  # 13 TType.MAP
       ('readContainerSet', 'writeContainerSet', True),  # 14 TType.SET
       ('readContainerList', 'writeContainerList', True),  # 15 TType.LIST
       (None, None, False),  # 16 TType.UTF8 # TODO: handle utf8 types?
       (None, None, False)  # 17 TType.UTF16 # TODO: handle utf16 types?
      )

  def _ttype_handlers(self, ttype, spec):
    if spec == 'BINARY':
      if ttype != TType.STRING:
        raise TProtocolException(type=TProtocolException.INVALID_DATA,
                                 message='Invalid binary field type %d' % ttype)
      return ('readBinary', 'writeBinary', False)
    if sys.version_info[0] == 2 and spec == 'UTF8':
      if ttype != TType.STRING:
        raise TProtocolException(type=TProtocolException.INVALID_DATA,
                                 message='Invalid string field type %d' % ttype)
      return ('readUtf8', 'writeUtf8', False)
    return self._TTYPE_HANDLERS[ttype] if ttype < len(self._TTYPE_HANDLERS) else (None, None, False)

  def _read_by_ttype(self, ttype, spec, espec):
    reader_name, _, is_container = self._ttype_handlers(ttype, spec)
    if reader_name is None:
      raise TProtocolException(type=TProtocolException.INVALID_DATA,
                               message='Invalid type %d' % (ttype))
    reader_func = getattr(self, reader_name)
    read = (lambda: reader_func(espec)) if is_container else reader_func
    while True:
      yield read()

  def readFieldByTType(self, ttype, spec):
    return self._read_by_ttype(ttype, spec, spec).next()

  def readContainerList(self, spec):
    ttype, tspec, is_immutable = spec
    (list_type, list_len) = self.readListBegin()
    # TODO: compare types we just decoded with thrift_spec
    elems = islice(self._read_by_ttype(ttype, spec, tspec), list_len)
    results = (tuple if is_immutable else list)(elems)
    self.readListEnd()
    return results

  def readContainerSet(self, spec):
    ttype, tspec, is_immutable = spec
    (set_type, set_len) = self.readSetBegin()
    # TODO: compare types we just decoded with thrift_spec
    elems = islice(self._read_by_ttype(ttype, spec, tspec), set_len)
    results = (frozenset if is_immutable else set)(elems)
    self.readSetEnd()
    return results

  def readContainerStruct(self, spec):
    (obj_class, obj_spec) = spec
    obj = obj_class()
    obj.read(self)
    return obj

  def readContainerMap(self, spec):
    ktype, kspec, vtype, vspec, is_immutable = spec
    (map_ktype, map_vtype, map_len) = self.readMapBegin()
    # TODO: compare types we just decoded with thrift_spec and
    # abort/skip if types disagree
    keys = self._read_by_ttype(ktype, spec, kspec)
    vals = self._read_by_ttype(vtype, spec, vspec)
    keyvals = islice(zip(keys, vals), map_len)
    results = (TFrozenDict if is_immutable else dict)(keyvals)
    self.readMapEnd()
    return results

  def readStruct(self, obj, thrift_spec, is_immutable=False):
    if is_immutable:
      fields = {}
    self.readStructBegin()
    while True:
      (fname, ftype, fid) = self.readFieldBegin()
      if ftype == TType.STOP:
        break
      try:
        field = thrift_spec[fid]
      except IndexError:
        self.skip(ftype)
      else:
        if field is not None and ftype == field[1]:
          fname = field[2]
          fspec = field[3]
          val = self.readFieldByTType(ftype, fspec)
          if is_immutable:
            fields[fname] = val
          else:
            setattr(obj, fname, val)
        else:
          self.skip(ftype)
      self.readFieldEnd()
    self.readStructEnd()
    if is_immutable:
      return obj(**fields)

  def writeContainerStruct(self, val, spec):
    val.write(self)

  def writeContainerList(self, val, spec):
    ttype, tspec, _ = spec
    self.writeListBegin(ttype, len(val))
    for _ in self._write_by_ttype(ttype, val, spec, tspec):
      pass
    self.writeListEnd()

  def writeContainerSet(self, val, spec):
    ttype, tspec, _ = spec
    self.writeSetBegin(ttype, len(val))
    for _ in self._write_by_ttype(ttype, val, spec, tspec):
      pass
    self.writeSetEnd()

  def writeContainerMap(self, val, spec):
    ktype, kspec, vtype, vspec, _ = spec
    self.writeMapBegin(ktype, vtype, len(val))
    for _ in zip(self._write_by_ttype(ktype, six.iterkeys(val), spec, kspec),
                 self._write_by_ttype(vtype, six.itervalues(val), spec, vspec)):
      pass
    self.writeMapEnd()

  def writeStruct(self, obj, thrift_spec):
    self.writeStructBegin(obj.__class__.__name__)
    for field in thrift_spec:
      if field is None:
        continue
      fname = field[2]
      val = getattr(obj, fname)
      if val is None:
        # skip writing out unset fields
        continue
      fid = field[0]
      ftype = field[1]
      fspec = field[3]
      self.writeFieldBegin(fname, ftype, fid)
      self.writeFieldByTType(ftype, val, fspec)
      self.writeFieldEnd()
    self.writeFieldStop()
    self.writeStructEnd()

  def _write_by_ttype(self, ttype, vals, spec, espec):
    _, writer_name, is_container = self._ttype_handlers(ttype, spec)
    writer_func = getattr(self, writer_name)
    write = (lambda v: writer_func(v, espec)) if is_container else writer_func
    for v in vals:
      yield write(v)

  def writeFieldByTType(self, ttype, val, spec):
    self._write_by_ttype(ttype, [val], spec, spec).next()


def checkIntegerLimits(i, bits):
    if bits == 8 and (i < -128 or i > 127):
        raise TProtocolException(TProtocolException.INVALID_DATA,
                                 "i8 requires -128 <= number <= 127")
    elif bits == 16 and (i < -32768 or i > 32767):
        raise TProtocolException(TProtocolException.INVALID_DATA,
                                 "i16 requires -32768 <= number <= 32767")
    elif bits == 32 and (i < -2147483648 or i > 2147483647):
        raise TProtocolException(TProtocolException.INVALID_DATA,
                                 "i32 requires -2147483648 <= number <= 2147483647")
    elif bits == 64 and (i < -9223372036854775808 or i > 9223372036854775807):
         raise TProtocolException(TProtocolException.INVALID_DATA,
                                  "i64 requires -9223372036854775808 <= number <= 9223372036854775807")


class TProtocolFactory(object):
  def getProtocol(self, trans):
    pass