def readFrame(self, req_sz): # the first word could either be the length field of a framed message # or the first bytes of an unframed message. first_word = self._transport.readAll(I32.size) frame_size, = I32.unpack(first_word) is_unframed = False if frame_size & TBinaryProtocol.VERSION_MASK == TBinaryProtocol.VERSION_1: self._set_client_type(THeaderClientType.UNFRAMED_BINARY) is_unframed = True elif (byte_index(first_word, 0) == TCompactProtocol.PROTOCOL_ID and byte_index(first_word, 1) & TCompactProtocol.VERSION_MASK == TCompactProtocol.VERSION): self._set_client_type(THeaderClientType.UNFRAMED_COMPACT) is_unframed = True if is_unframed: bytes_left_to_read = req_sz - I32.size if bytes_left_to_read > 0: rest = self._transport.read(bytes_left_to_read) else: rest = b"" self._read_buffer = BufferIO(first_word + rest) return # ok, we're still here so we're framed. if frame_size > self._max_frame_size: raise TTransportException( TTransportException.SIZE_LIMIT, "Frame was too large.", ) read_buffer = BufferIO(self._transport.readAll(frame_size)) # the next word is either going to be the version field of a # binary/compact protocol message or the magic value + flags of a # header protocol message. second_word = read_buffer.read(I32.size) version, = I32.unpack(second_word) read_buffer.seek(0) if version >> 16 == HEADER_MAGIC: self._set_client_type(THeaderClientType.HEADERS) self._read_buffer = self._parse_header_format(read_buffer) elif version & TBinaryProtocol.VERSION_MASK == TBinaryProtocol.VERSION_1: self._set_client_type(THeaderClientType.FRAMED_BINARY) self._read_buffer = read_buffer elif (byte_index(second_word, 0) == TCompactProtocol.PROTOCOL_ID and byte_index(second_word, 1) & TCompactProtocol.VERSION_MASK == TCompactProtocol.VERSION): self._set_client_type(THeaderClientType.FRAMED_COMPACT) self._read_buffer = read_buffer else: raise TTransportException( TTransportException.INVALID_CLIENT_TYPE, "Could not detect client transport type.", )
class TTransformTransport(TTransport.TTransportBase): """ A transport that transforms the written payload and supplies it as read payload. The transformation can be any kind of Synchronous operation, either a computation or a result of a network call. """ def __init__(self, value=None, offset=0): """ :param value: a value to read from for stringio If value is set, the read buffer will be initialized with it otherwise, it is for writing :param offset: the offset to start reading from. """ if value is not None: self.__read_buffer = BufferIO(value) else: self.__read_buffer = BufferIO() if offset: self.__read_buffer.seek(offset) self.__write_buffer = BufferIO() def isOpen(self): """ :return: True if the transport is open, False otherwise """ return (not self.__read_buffer.closed) and ( not self.__write_buffer.closed) def close(self): """ closes the transport """ self.__read_buffer.close() self.__write_buffer.close() def read(self, sz): """ reads from the transport :param sz: number of bytes to read :return: the data read """ return self.__read_buffer.read(sz) def write(self, buf): """ writes to the transport :param buf: the buffer to write """ self.__write_buffer.write(buf) def _transform(self, buf): """ Transforms the data written, and sets it as data to be read :param buf: The data written to the transport :return: The data to set as readable from the transport """ return buf def flush(self): """ flushes the transport (verify all previous writes actually written) """ self.__read_buffer = BufferIO( self._transform(self.__write_buffer.getvalue())) self.__write_buffer = BufferIO() def getvalue(self): """ :return: all the current data available for read from the transport """ return self.__read_buffer.getvalue()