def test_basic(self):
buf = ByteBuffer(4)
assert buf.getlength() == 4
assert buf.getitem(2) == '\x00'
buf.setitem(0, 'A')
buf.setitem(3, 'Z')
assert buf.as_str() == 'A\x00\x00Z'
def read_console_wide(space, handle, maxlen):
"""
Make a blocking call to ReadConsoleW, returns wchar-encoded bytes
"""
err = 0
sig = 0
# Windows uses a 16-bit wchar_t, we mimic that with bytes
buf = ByteBuffer((maxlen + 2) * 2)
addr = buf.get_raw_address()
off = 0 # offset from the beginning of buf, in wchar
# readlen is in 16 bits, readlen_b is in 8-bit bytes
readlen = readlen_b = 0
bufsize = BUFSIZ
with lltype.scoped_alloc(rwin32.LPDWORD.TO, 1) as n:
while readlen_b < maxlen:
neg_one = rffi.cast(rwin32.DWORD, -1)
n[0] = neg_one
length = min(maxlen - off, bufsize)
rwin32.SetLastError_saved(0)
res = rwin32.ReadConsoleW(
handle, rffi.cast(rwin32.LPWSTR, rffi.ptradd(addr, readlen_b)),
length, n, rffi.cast(rwin32.LPVOID, 0))
nread = intmask(n[0])
err = rwin32.GetLastError_saved()
if not res:
break
if nread == -1 and err == rwin32.ERROR_OPERATION_ABORTED:
break
if nread == 0:
if err != rwin32.ERROR_OPERATION_ABORTED:
break
err = 0
# This will only catch CTRL-C on the main thread
sleep(space, space.newfloat(0.1))
continue
readlen += nread
readlen_b = 2 * readlen
# We didn't manage to read the whole buffer
# don't try again as it will just block
if nread < length:
break
if buf.getitem(readlen_b - 2) == '\n':
# We read a new line
break
# If the buffer ends with a high surrogate, take an extra character.
if (readlen_b + 1) >= maxlen:
with lltype.scoped_alloc(rwin32.LPWORD.TO, 1) as char_type:
ptr = rffi.cast(rffi.CWCHARP, rffi.ptradd(addr, +1))
rwin32.GetStringTypeW(rwin32.CT_CTYPE3, ptr, 1, char_type)
if intmask(char_type[0]) == intmask(
rwin32.C3_HIGHSURROGATE):
readlen_b += 2
break
if err:
raise OperationError(space.w_WindowsError, space.newint(err))
if readlen_b <= 0 or buf.getitem(0) == '\x1a':
return ''
else:
return buf.getslice(0, 1, readlen_b)
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!
class BufferedMixin:
_mixin_ = True
def __init__(self, space):
W_IOBase.__init__(self, space)
self.state = STATE_ZERO
self.buffer = None
self.abs_pos = 0 # Absolute position inside the raw stream (-1 if
# unknown).
self.pos = 0 # Current logical position in the buffer
self.raw_pos = 0 # Position of the raw stream in the buffer.
self.read_end = -1 # Just after the last buffered byte in the buffer,
# or -1 if the buffer isn't ready for reading
self.write_pos = 0 # Just after the last byte actually written
self.write_end = -1 # Just after the last byte waiting to be written,
# or -1 if the buffer isn't ready for writing.
self.lock = None
self.readable = False
self.writable = False
def _reader_reset_buf(self):
self.read_end = -1
def _writer_reset_buf(self):
self.write_pos = 0
self.write_end = -1
def _init(self, space):
if self.buffer_size <= 0:
raise oefmt(space.w_ValueError,
"buffer size must be strictly positive")
if space.config.translation.split_gc_address_space:
# When using split GC address space, it is not possible to get the
# raw address of a GC buffer. Therefore we use a buffer backed by
# raw memory.
self.buffer = RawByteBuffer(self.buffer_size)
else:
# TODO: test whether using the raw buffer is faster
self.buffer = ByteBuffer(self.buffer_size)
self.lock = TryLock(space)
try:
self._raw_tell(space)
except OperationError:
pass
def _check_init(self, space):
if self.state == STATE_ZERO:
raise oefmt(space.w_ValueError,
"I/O operation on uninitialized object")
elif self.state == STATE_DETACHED:
raise oefmt(space.w_ValueError, "raw stream has been detached")
def _check_closed(self, space, message=None):
self._check_init(space)
W_IOBase._check_closed(self, space, message)
def _raw_tell(self, space):
w_pos = space.call_method(self.w_raw, "tell")
pos = space.r_longlong_w(w_pos)
if pos < 0:
raise oefmt(space.w_IOError,
"raw stream returned invalid position")
self.abs_pos = pos
return pos
def closed_get_w(self, space):
self._check_init(space)
return space.getattr(self.w_raw, space.newtext("closed"))
def name_get_w(self, space):
self._check_init(space)
return space.getattr(self.w_raw, space.newtext("name"))
def mode_get_w(self, space):
self._check_init(space)
return space.getattr(self.w_raw, space.newtext("mode"))
def readable_w(self, space):
self._check_init(space)
return space.call_method(self.w_raw, "readable")
def writable_w(self, space):
self._check_init(space)
return space.call_method(self.w_raw, "writable")
def seekable_w(self, space):
self._check_init(space)
return space.call_method(self.w_raw, "seekable")
def isatty_w(self, space):
self._check_init(space)
return space.call_method(self.w_raw, "isatty")
def repr_w(self, space):
typename = space.type(self).name
try:
w_name = space.getattr(self, space.newtext("name"))
except OperationError as e:
if not e.match(space, space.w_Exception):
raise
return space.newtext("<%s>" % (typename, ))
else:
name_repr = space.text_w(space.repr(w_name))
return space.newtext("<%s name=%s>" % (typename, name_repr))
# ______________________________________________
@signature(types.any(), returns=types.int())
def _readahead(self):
if self.readable and self.read_end != -1:
available = self.read_end - self.pos
assert available >= 0
return available
return 0
def _raw_offset(self):
if self.raw_pos >= 0 and ((self.readable and self.read_end != -1) or
(self.writable and self.write_end != -1)):
return self.raw_pos - self.pos
return 0
def tell_w(self, space):
self._check_init(space)
pos = self._raw_tell(space) - self._raw_offset()
return space.newint(pos)
@unwrap_spec(pos=r_longlong, whence=int)
def seek_w(self, space, pos, whence=0):
self._check_init(space)
if whence not in (0, 1, 2):
raise oefmt(space.w_ValueError,
"whence must be between 0 and 2, not %d", whence)
self._check_closed(space, "seek of closed file")
if whence != 2 and self.readable:
# Check if seeking leaves us inside the current buffer, so as to
# return quickly if possible. Also, we needn't take the lock in
# this fast path.
if self.abs_pos == -1:
self._raw_tell(space)
current = self.abs_pos
available = self._readahead()
if available > 0:
if whence == 0:
offset = pos - (current - self._raw_offset())
else:
offset = pos
if -self.pos <= offset <= available:
newpos = self.pos + int(offset)
assert newpos >= 0
self.pos = newpos
return space.newint(current - available + offset)
# Fallback: invoke raw seek() method and clear buffer
with self.lock:
if self.writable:
self._writer_flush_unlocked(space)
self._writer_reset_buf()
if whence == 1:
pos -= self._raw_offset()
n = self._raw_seek(space, pos, whence)
self.raw_pos = -1
if self.readable:
self._reader_reset_buf()
return space.newint(n)
def _raw_seek(self, space, pos, whence):
w_pos = space.call_method(self.w_raw, "seek", space.newint(pos),
space.newint(whence))
pos = space.r_longlong_w(w_pos)
if pos < 0:
raise oefmt(space.w_IOError,
"Raw stream returned invalid position")
self.abs_pos = pos
return pos
def _closed(self, space):
return space.is_true(space.getattr(self.w_raw,
space.newtext("closed")))
def close_w(self, space):
self._check_init(space)
with self.lock:
if self._closed(space):
return
try:
space.call_method(self, "flush")
finally:
with self.lock:
space.call_method(self.w_raw, "close")
def simple_flush_w(self, space):
self._check_init(space)
return space.call_method(self.w_raw, "flush")
def _writer_flush_unlocked(self, space):
if self.write_end == -1 or self.write_pos == self.write_end:
return
# First, rewind
rewind = self._raw_offset() + (self.pos - self.write_pos)
if rewind != 0:
self._raw_seek(space, -rewind, 1)
self.raw_pos -= rewind
written = 0
while self.write_pos < self.write_end:
try:
n = self._raw_write(space, self.write_pos, self.write_end)
except BlockingIOError:
raise make_write_blocking_error(space, 0)
self.write_pos += n
self.raw_pos = self.write_pos
written += n
# Partial writes can return successfully when interrupted by a
# signal (see write(2)). We must run signal handlers before
# blocking another time, possibly indefinitely.
space.getexecutioncontext().checksignals()
self._writer_reset_buf()
def _write(self, space, data):
w_data = space.newbytes(data)
while True:
try:
w_written = space.call_method(self.w_raw, "write", w_data)
except OperationError as e:
if trap_eintr(space, e):
continue # try again
raise
else:
break
if space.is_w(w_written, space.w_None):
# Non-blocking stream would have blocked.
raise BlockingIOError()
written = space.getindex_w(w_written, space.w_IOError)
if not 0 <= written <= len(data):
raise oefmt(space.w_IOError, "raw write() returned invalid length")
if self.abs_pos != -1:
self.abs_pos += written
return written
def _raw_write(self, space, start, end):
return self._write(space, self.buffer[start:end])
def detach_w(self, space):
self._check_init(space)
space.call_method(self, "flush")
w_raw = self.w_raw
self.w_raw = None
self.state = STATE_DETACHED
return w_raw
def fileno_w(self, space):
self._check_init(space)
return space.call_method(self.w_raw, "fileno")
@unwrap_spec(w_size=WrappedDefault(None))
def truncate_w(self, space, w_size):
self._check_init(space)
with self.lock:
if self.writable:
self._flush_and_rewind_unlocked(space)
# invalidate cached position
self.abs_pos = -1
return space.call_method(self.w_raw, "truncate", w_size)
# ________________________________________________________________
# Read methods
def read_w(self, space, w_size=None):
self._check_init(space)
self._check_closed(space, "read of closed file")
size = convert_size(space, w_size)
if size == -1:
# read until the end of stream
with self.lock:
return self._read_all(space)
elif size >= 0:
res = self._read_fast(size)
if res is None:
with self.lock:
res = self._read_generic(space, size)
else:
raise oefmt(space.w_ValueError,
"read length must be positive or -1")
return space.newbytes(res)
@unwrap_spec(size=int)
def peek_w(self, space, size=0):
self._check_init(space)
self._check_closed(space, "peek of closed file")
with self.lock:
if self.writable:
self._flush_and_rewind_unlocked(space)
# Constraints:
# 1. we don't want to advance the file position.
# 2. we don't want to lose block alignment, so we can't shift the
# buffer to make some place.
# Therefore, we either return `have` bytes (if > 0), or a full
# buffer.
have = self._readahead()
if have > 0:
data = self.buffer[self.pos:self.pos + have]
return space.newbytes(data)
# Fill the buffer from the raw stream, and copy it to the result
self._reader_reset_buf()
try:
size = self._fill_buffer(space)
except BlockingIOError:
size = 0
self.pos = 0
data = self.buffer[0:size]
return space.newbytes(data)
@unwrap_spec(size=int)
def read1_w(self, space, size):
self._check_init(space)
self._check_closed(space, "read of closed file")
if size < 0:
raise oefmt(space.w_ValueError, "read length must be positive")
if size == 0:
return space.newbytes("")
with self.lock:
# Return up to n bytes. If at least one byte is buffered, we only
# return buffered bytes. Otherwise, we do one raw read.
# XXX: this mimicks the io.py implementation but is probably
# wrong. If we need to read from the raw stream, then we could
# actually read all `n` bytes asked by the caller (and possibly
# more, so as to fill our buffer for the next reads).
have = self._readahead()
if have == 0:
if self.writable:
self._flush_and_rewind_unlocked(space)
# Fill the buffer from the raw stream
self._reader_reset_buf()
self.pos = 0
try:
have = self._fill_buffer(space)
except BlockingIOError:
have = 0
if size > have:
size = have
endpos = self.pos + size
data = self.buffer[self.pos:endpos]
self.pos = endpos
return space.newbytes(data)
def _read_all(self, space):
"Read all the file, don't update the cache"
# Must run with the lock held!
builder = StringBuilder()
# First copy what we have in the current buffer
current_size = self._readahead()
data = None
if current_size:
data = self.buffer[self.pos:self.pos + current_size]
builder.append(data)
self.pos += current_size
# We're going past the buffer's bounds, flush it
if self.writable:
self._flush_and_rewind_unlocked(space)
self._reader_reset_buf()
while True:
# Read until EOF or until read() would block
w_data = space.call_method(self.w_raw, "read")
if space.is_w(w_data, space.w_None):
if current_size == 0:
return w_data
break
data = space.bytes_w(w_data)
size = len(data)
if size == 0:
break
builder.append(data)
current_size += size
if self.abs_pos != -1:
self.abs_pos += size
return space.newbytes(builder.build())
def _raw_read(self, space, buffer, start, length):
assert buffer is not None
length = intmask(length)
start = intmask(start)
w_view = SimpleView(SubBuffer(buffer, start, length)).wrap(space)
while True:
try:
w_size = space.call_method(self.w_raw, "readinto", w_view)
except OperationError as e:
if trap_eintr(space, e):
continue # try again
raise
else:
break
if space.is_w(w_size, space.w_None):
raise BlockingIOError()
size = space.int_w(w_size)
if size < 0 or size > length:
raise oefmt(
space.w_IOError,
"raw readinto() returned invalid length %d (should "
"have been between 0 and %d)", size, length)
if self.abs_pos != -1:
self.abs_pos += size
return size
def _fill_buffer(self, space):
start = self.read_end
if start == -1:
start = 0
length = self.buffer_size - start
size = self._raw_read(space, self.buffer, start, length)
if size > 0:
self.read_end = self.raw_pos = start + size
return size
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:
self.output_slice(space, result_buffer, 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
self.output_slice(space, result_buffer, written,
self.buffer[self.pos:self.pos + size])
self.pos += size
written += size
remaining -= size
return result_buffer[0:written]
def _read_fast(self, n):
"""Read n bytes from the buffer if it can, otherwise return None.
This function is simple enough that it can run unlocked."""
current_size = self._readahead()
if n <= current_size:
endpos = self.pos + n
res = self.buffer[self.pos:endpos]
self.pos = endpos
return res
return None
def readline_w(self, space, w_limit=None):
self._check_init(space)
self._check_closed(space, "readline of closed file")
limit = convert_size(space, w_limit)
# First, try to find a line in the buffer. This can run
# unlocked because the calls to the C API are simple enough
# that they can't trigger any thread switch.
have = self._readahead()
if limit >= 0 and have > limit:
have = limit
for pos in range(self.pos, self.pos + have):
if self.buffer[pos] == '\n':
break
else:
pos = -1
if pos >= 0:
w_res = space.newbytes(self.buffer[self.pos:pos + 1])
self.pos = pos + 1
return w_res
if have == limit:
w_res = space.newbytes(self.buffer[self.pos:self.pos + have])
self.pos += have
return w_res
written = 0
with self.lock:
# Now we try to get some more from the raw stream
chunks = []
if have > 0:
chunks.append(self.buffer[self.pos:self.pos + have])
written += have
self.pos += have
if limit >= 0:
limit -= have
if self.writable:
self._flush_and_rewind_unlocked(space)
while True:
self._reader_reset_buf()
have = self._fill_buffer(space)
if have == 0:
break
if limit >= 0 and have > limit:
have = limit
pos = 0
found = False
while pos < have:
c = self.buffer.getitem(pos)
pos += 1
if c == '\n':
self.pos = pos
found = True
break
chunks.append(self.buffer[0:pos])
if found:
break
if have == limit:
self.pos = have
break
written += have
if limit >= 0:
limit -= have
return space.newbytes(''.join(chunks))
# ____________________________________________________
# Write methods
def _adjust_position(self, new_pos):
assert new_pos >= 0
self.pos = new_pos
if self.readable and self.read_end != -1 and self.read_end < new_pos:
self.read_end = self.pos
def write_w(self, space, w_data):
self._check_init(space)
self._check_closed(space, "write to closed file")
data = space.getarg_w('s*', w_data).as_str()
size = len(data)
with self.lock:
if (not (self.readable and self.read_end != -1)
and not (self.writable and self.write_end != -1)):
self.pos = 0
self.raw_pos = 0
available = self.buffer_size - self.pos
# Fast path: the data to write can be fully buffered
if size <= available:
for i in range(size):
self.buffer[self.pos + i] = data[i]
if self.write_end == -1 or self.write_pos > self.pos:
self.write_pos = self.pos
self._adjust_position(self.pos + size)
if self.pos > self.write_end:
self.write_end = self.pos
return space.newint(size)
# First write the current buffer
try:
self._writer_flush_unlocked(space)
except OperationError as e:
if not e.match(space,
space.gettypeobject(W_BlockingIOError.typedef)):
raise
w_exc = e.get_w_value(space)
assert isinstance(w_exc, W_BlockingIOError)
if self.readable:
self._reader_reset_buf()
# Make some place by shifting the buffer
for i in range(self.write_pos, self.write_end):
self.buffer.setitem(i - self.write_pos,
self.buffer.getitem(i))
self.write_end -= self.write_pos
self.raw_pos -= self.write_pos
newpos = self.pos - self.write_pos
assert newpos >= 0
self.pos = newpos
self.write_pos = 0
available = self.buffer_size - self.write_end
assert available >= 0
if size <= available:
# Everything can be buffered
for i in range(size):
self.buffer[self.write_end + i] = data[i]
self.write_end += size
self.pos += size
return space.newint(size)
# Buffer as much as possible
for i in range(available):
self.buffer[self.write_end + i] = data[i]
self.write_end += available
self.pos += available
# Modifying the existing exception will will change
# e.characters_written but not e.args[2]. Therefore
# we just replace with a new error.
raise make_write_blocking_error(space, available)
# Adjust the raw stream position if it is away from the logical
# stream position. This happens if the read buffer has been filled
# but not modified (and therefore _bufferedwriter_flush_unlocked()
# didn't rewind the raw stream by itself).
offset = self._raw_offset()
if offset:
self._raw_seek(space, -offset, 1)
self.raw_pos -= offset
# Then write buf itself. At this point the buffer has been emptied
remaining = size
written = 0
while remaining > self.buffer_size:
try:
n = self._write(space, data[written:])
except BlockingIOError:
# Write failed because raw file is non-blocking
if remaining > self.buffer_size:
# Can't buffer everything, still buffer as much as
# possible
for i in range(self.buffer_size):
self.buffer[i] = data[written + i]
self.raw_pos = 0
self._adjust_position(self.buffer_size)
self.write_end = self.buffer_size
written += self.buffer_size
raise make_write_blocking_error(space, written)
break
written += n
remaining -= n
# Partial writes can return successfully when interrupted by a
# signal (see write(2)). We must run signal handlers before
# blocking another time, possibly indefinitely.
space.getexecutioncontext().checksignals()
if self.readable:
self._reader_reset_buf()
if remaining > 0:
for i in range(remaining):
self.buffer[i] = data[written + i]
written += remaining
self.write_pos = 0
self.write_end = remaining
self._adjust_position(remaining)
self.raw_pos = 0
return space.newint(written)
def flush_w(self, space):
self._check_init(space)
self._check_closed(space, "flush of closed file")
with self.lock:
self._flush_and_rewind_unlocked(space)
def _flush_and_rewind_unlocked(self, space):
self._writer_flush_unlocked(space)
if self.readable:
# Rewind the raw stream so that its position corresponds to
# the current logical position.
try:
self._raw_seek(space, -self._raw_offset(), 1)
finally:
self._reader_reset_buf()