def make_buffer(flag):
if flag:
buf = ByteBuffer(len(DATA))
buf.setslice(0, DATA)
else:
buf = StringBuffer(DATA)
return buf
def _read_generic(self, space, n):
"""Generic read function: read from the stream until enough bytes are
read, or until an EOF occurs or until read() would block."""
# Must run with the lock held!
current_size = self._readahead()
if n <= current_size:
return self._read_fast(n)
result_buffer = ByteBuffer(n)
remaining = n
written = 0
if current_size:
result_buffer.setslice(
written, self.buffer[self.pos:self.pos + current_size])
remaining -= current_size
written += current_size
self.pos += current_size
# Flush the write buffer if necessary
if self.writable:
self._flush_and_rewind_unlocked(space)
self._reader_reset_buf()
# Read whole blocks, and don't buffer them
while remaining > 0:
r = self.buffer_size * (remaining // self.buffer_size)
if r == 0:
break
try:
size = self._raw_read(space, result_buffer, written, r)
except BlockingIOError:
if written == 0:
return None
size = 0
if size == 0:
return result_buffer[0:written]
remaining -= size
written += size
self.pos = 0
self.raw_pos = 0
self.read_end = 0
while remaining > 0 and self.read_end < self.buffer_size:
try:
size = self._fill_buffer(space)
except BlockingIOError:
# EOF or read() would block
if written == 0:
return None
size = 0
if size == 0:
break
if remaining > 0:
if size > remaining:
size = remaining
result_buffer.setslice(written,
self.buffer[self.pos:self.pos + size])
self.pos += size
written += size
remaining -= size
return result_buffer[0:written]
def test_getslice_shortcut(self):
buf = ByteBuffer(4)
buf.setslice(0, b"data")
buf.getitem = None
assert buf.getslice(0, 2, 1, 2) == b"da" # no crash!
def test_typed_read(self):
data = struct.pack('HH', 0x1234, 0x5678)
buf = ByteBuffer(4)
buf.setslice(0, data)
assert buf.typed_read(rffi.USHORT, 0) == 0x1234
assert buf.typed_read(rffi.USHORT, 2) == 0x5678