def test_pipe_iostream_big_write(self):
r, w = os.pipe()
rs = PipeIOStream(r, io_loop=self.io_loop)
ws = PipeIOStream(w, io_loop=self.io_loop)
NUM_BYTES = 1048576
# Write 1MB of data, which should fill the buffer
ws.write(b"1" * NUM_BYTES)
rs.read_bytes(NUM_BYTES, self.stop)
data = self.wait()
self.assertEqual(data, b"1" * NUM_BYTES)
ws.close()
rs.close()
def test_pipe_iostream_big_write(self):
r, w = os.pipe()
rs = PipeIOStream(r, io_loop=self.io_loop)
ws = PipeIOStream(w, io_loop=self.io_loop)
NUM_BYTES = 1048576
# Write 1MB of data, which should fill the buffer
ws.write(b"1" * NUM_BYTES)
rs.read_bytes(NUM_BYTES, self.stop)
data = self.wait()
self.assertEqual(data, b"1" * NUM_BYTES)
ws.close()
rs.close()
class SafeFileIOStream:
def __init__(self, fname):
self.fname = fname
files_dir = os.path.dirname(fname)
if not os.path.exists(files_dir):
os.mkdir(files_dir)
def __enter__(self):
# Create file
os.open(self.fname, os.O_CREAT)
# Create stream
fd = os.open(self.fname, os.O_WRONLY)
self.stream = PipeIOStream(fd)
return self.stream
def __exit__(self, exc_type, exc_val, exc_tb):
# Close stream
self.stream.close()
def test_pipe_iostream(self):
r, w = os.pipe()
rs = PipeIOStream(r, io_loop=self.io_loop)
ws = PipeIOStream(w, io_loop=self.io_loop)
ws.write(b"hel")
ws.write(b"lo world")
rs.read_until(b' ', callback=self.stop)
data = self.wait()
self.assertEqual(data, b"hello ")
rs.read_bytes(3, self.stop)
data = self.wait()
self.assertEqual(data, b"wor")
ws.close()
rs.read_until_close(self.stop)
data = self.wait()
self.assertEqual(data, b"ld")
rs.close()
def test_pipe_iostream(self):
r, w = os.pipe()
rs = PipeIOStream(r, io_loop=self.io_loop)
ws = PipeIOStream(w, io_loop=self.io_loop)
ws.write(b"hel")
ws.write(b"lo world")
rs.read_until(b" ", callback=self.stop)
data = self.wait()
self.assertEqual(data, b"hello ")
rs.read_bytes(3, self.stop)
data = self.wait()
self.assertEqual(data, b"wor")
ws.close()
rs.read_until_close(self.stop)
data = self.wait()
self.assertEqual(data, b"ld")
rs.close()
class PipeStream(Stream):
def __init__(self, rpipe, wpipe=None, auto_close=False):
"""Pipe-based stream
NOTE: reading from or writing to files, use os.open to get the file
descriptor instead of python's open. Socket file descriptors and
others are fine.
when you use os.pipe to generate one write pipe and one read pipe, you
need to pass both of them into init method.
:param rpipe: an integer file descriptor which supports read ops
:param wpipe: an integer file descriptor which supports write ops
:param auto: flag to indicate to close the stream automatically or not
"""
assert rpipe is not None
self._rpipe = rpipe
self._wpipe = wpipe
self._rs = PipeIOStream(self._rpipe) if self._rpipe is not None else None
self._ws = PipeIOStream(self._wpipe) if self._wpipe is not None else None
self.auto_close = auto_close
self.state = StreamState.init
self.exception = None
@tornado.gen.coroutine
def read(self):
if self.exception:
raise self.exception
if self.state == StreamState.completed or self._rpipe is None:
raise tornado.gen.Return("")
elif self.state == StreamState.init:
self.state = StreamState.streaming
chunk = ""
try:
chunk = yield self._rs.read_bytes(common.MAX_PAYLOAD_SIZE, partial=True)
except StreamClosedError:
# reach the end of the pipe stream
self.state = StreamState.completed
finally:
if self.exception:
raise self.exception
raise tornado.gen.Return(chunk)
@tornado.gen.coroutine
def write(self, chunk):
assert self._wpipe is not None
if self.exception:
raise self.exception
try:
yield self._ws.write(chunk)
self.state = StreamState.streaming
except StreamClosedError:
self.state = StreamState.completed
raise UnexpectedError("Stream has been closed.")
finally:
if self.exception:
raise self.exception
def set_exception(self, exception):
self.exception = exception
self.close()
def close(self):
self.state = StreamState.completed
if self._ws and self.auto_close:
self._ws.close()
if self._rs and self.auto_close:
self._rs.close()
