class TestingPipe(Pipe): # reads come out in the same blocks as the writes # ...wraps LocalPipe def __init__(self, layout: str, name: str = None, stream=None): super().__init__(name=name, layout=layout, stream=stream) self._pipe = LocalPipe(layout, name, stream) self._closed = False self.data_blocks: deque = deque() async def write(self, data): if self._closed: raise PipeError("Cannot write to a closed pipe") self.data_blocks.append(data) def write_nowait(self, data): if self._closed: raise PipeError("Cannot write to a closed pipe") self.data_blocks.append(data) def close_interval_no_wait(self): self.data_blocks.append(None) async def close_interval(self): self.data_blocks.append(None) def consume(self, rows): return self._pipe.consume(rows) def reread_last(self): self._pipe.reread_last() @property def end_of_interval(self): return self._pipe.end_of_interval async def read(self, flatten=False): # first write a block to the pipe if there are any waiting if len(self.data_blocks) > 0: block = self.data_blocks.popleft() if block is None: await self._pipe.close_interval() else: await self._pipe.write(block) elif self._closed: await self._pipe.close() # now return the result of the inner pipe's read return await self._pipe.read(flatten) def is_empty(self): return self._pipe.is_empty() async def close(self): self._closed = True
async def _run(): LAYOUT = "float32_3" LENGTH = 100 pipe = LocalPipe(layout="float32_3") tx_data = helpers.create_data(LAYOUT, length=LENGTH) # write a block and close off the pipe await pipe.write(tx_data) await pipe.close() # read it back rx_data = await pipe.read() np.testing.assert_array_equal(tx_data, rx_data) # don't consume it so you can read it agai pipe.consume(len(rx_data[:500])) # read it back again rx_data = await pipe.read() np.testing.assert_array_equal(tx_data[500:], rx_data) # the pipe should be empty now self.assertTrue(pipe.is_empty()) with self.assertRaises(EmptyPipeError): await pipe.read()