tcpclient_write_tornado_3.2.2.py

import logging
import tornado
import tornado.gen

from tornado.iostream import IOStream
from tornado.ioloop import IOLoop
from tornado.concurrent import Future
import socket


class TCPClient(object):

    def __init__(self, host, port):
        self.host = host
        self.port = int(port)
        self._iostream = None

    def _get_iostream(self):
        future = IOStreamFuture()
        if self._iostream is None or self._iostream.closed():
            ioloop = IOLoop.current()
            self._iostream = IOStream(socket.socket(socket.AF_INET), io_loop=ioloop)
            future._iostream = self._iostream
            self._iostream.set_close_callback(future.resolve)
            logging.info('Connecting to %s %s', self.host, self.port)
            self._iostream.connect((self.host, self.port), callback=future.resolve)
        else:
            future.set_result(self._iostream)
        return future

    @tornado.gen.coroutine
    def write(self, data):
        try:
            if not isinstance(data, bytes):
                data = data.encode('utf-8')
            conn = yield self._get_iostream()
            logging.info('connected')
            yield tornado.gen.Task(conn.write, data)
        except Exception as exc:
            logging.exception('TCPClient error %s', exc)
            self._iostream = None
            raise


class IOStreamFuture(Future):
    def resolve(self):
        if self._iostream is None or self._iostream.closed():
            self.set_exception(IOError('IOStream is closed'))
        else:
            self.set_result(self._iostream)