from rpython.rlib.rutf8 import ( codepoints_in_utf8, codepoint_at_pos, Utf8StringIterator, Utf8StringBuilder) from rpython.rlib.signature import signature, finishsigs from rpython.rlib import types from pypy.interpreter.error import OperationError, oefmt from pypy.interpreter.typedef import ( TypeDef, generic_new_descr, GetSetProperty) from pypy.interpreter.gateway import interp2app, unwrap_spec, WrappedDefault from pypy.module._io.interp_textio import ( W_TextIOBase, W_IncrementalNewlineDecoder) from pypy.module._io.interp_iobase import convert_size from pypy.objspace.std.unicodeobject import W_UnicodeObject from rpython.rlib import rarithmetic def _find_end(start, size, total): available = total - start if size >= 0 and size <= available: end = start + size else: end = total assert 0 <= start and 0 <= end return end class UnicodeIO(object): def __init__(self, data=None, length=0): if data is None: data = '' self.data = [] # list of r_int32 self.write(data, length, 0) def getdata(self, i): return rarithmetic.widen(self.data[i]) def setdata(self, i, val): self.data[i] = rarithmetic.r_int32(val) def getdata_slice(self, start, end): builder = Utf8StringBuilder(end - start) for index in range(start, end): builder.append_code(self.getdata(index)) return builder def resize(self, newlength): if len(self.data) > newlength: self.data = self.data[:newlength] if len(self.data) < newlength: self.data.extend([0] * (newlength - len(self.data))) def read(self, start, size): end = _find_end(start, size, len(self.data)) return self.getdata_slice(start, end) def _convert_limit(self, limit, start): if limit < 0 or limit > len(self.data) - start: limit = len(self.data) - start if limit < 0: # happens when self.pos > len(self.data) limit = 0 return limit def readline_universal(self, start, limit): # Universal newline search. Find any of \r, \r\n, \n limit = self._convert_limit(limit, start) end = start + limit pos = start builder = Utf8StringBuilder() while pos < end: ch = self.getdata(pos) builder.append_code(ch) pos += 1 if ch == ord('\n'): break if ch == ord('\r'): if pos >= end: break if self.getdata(pos) == ord('\n'): builder.append_code(ch) pos += 1 break else: break return builder def readline(self, marker, start, limit): limit = self._convert_limit(limit, start) end = start + limit found = False for pos in range(start, end - len(marker) + 1): ch = self.getdata(pos) if ch == ord(marker[0]): for j in range(1, len(marker)): if self.getdata(pos + j) != ord(marker[j]): break # from inner loop else: pos += len(marker) found = True break if not found: pos = end return self.getdata_slice(start, pos) def write(self, string, length, start): newlen = start + length if newlen > len(self.data): self.resize(newlen) for ch in Utf8StringIterator(string): self.setdata(start, ch) start += 1 return length def truncate(self, size): if size < len(self.data): self.resize(size) def getvalue(self): return self.getdata_slice(0, len(self.data)) READING, ACCUMULATING, RWBUFFER, CLOSED = range(4) @finishsigs class W_StringIO(W_TextIOBase): def __init__(self, space): W_TextIOBase.__init__(self, space) self.buf = None self.w_value = W_UnicodeObject.EMPTY self.builder = None self.pos = 0 self.state = READING def get_length(self): """Return the total size (in codepoints) of the object""" if self.state == READING: return self.w_value._len() elif self.state == ACCUMULATING: return self.builder.getlength() else: return len(self.buf.data) def _init_newline(self, space, w_newline): self.writenl = None if space.is_w(w_newline, space.w_None): newline = None elif space.isinstance_w(w_newline, space.w_unicode): newline = space.utf8_w(w_newline) else: raise oefmt(space.w_TypeError, "newline must be str or None, not %T", w_newline) if (newline is not None and newline != "" and newline != "\n" and newline != "\r" and newline != "\r\n"): # Not using oefmt() because I don't know how to use it # with unicode raise OperationError(space.w_ValueError, space.mod( space.newtext("illegal newline value: %s"), w_newline ) ) self.w_decoder = None self.readnl = None if newline is not None: self.readnl = newline self.readuniversal = newline is None or newline == "" self.readtranslate = newline is None if newline and newline[0] == "\r": self.writenl = newline if self.readuniversal: self.w_decoder = space.call_function( space.gettypefor(W_IncrementalNewlineDecoder), space.w_None, space.newint(int(self.readtranslate)) ) @unwrap_spec(w_newline = WrappedDefault(u"\n")) def descr_init(self, space, w_initvalue=None, w_newline=None): self._init_newline(space, w_newline) if not space.is_none(w_initvalue): self.w_value = self._decode_string(space, w_initvalue) else: self.w_value = W_UnicodeObject.EMPTY # In case __init__ is called multiple times self.buf = None self.pos = 0 self.state = READING def descr_getstate(self, space): w_initialval = self.getvalue_w(space) w_dict = space.call_method(self.getdict(space), "copy") readnl = self.readnl if readnl is None: w_readnl = space.w_None else: w_readnl = space.str(space.newutf8(readnl, codepoints_in_utf8(readnl))) # YYY return space.newtuple([ w_initialval, w_readnl, space.newint(self.pos), w_dict ]) def descr_setstate(self, space, w_state): self._check_closed(space) # We allow the state tuple to be longer than 4, because we may need # someday to extend the object's state without breaking # backwards-compatibility if (not space.isinstance_w(w_state, space.w_tuple) or space.len_w(w_state) < 4): raise oefmt(space.w_TypeError, "%T.__setstate__ argument should be a 4-tuple, got %T", self, w_state) w_initval, w_readnl, w_pos, w_dict = space.unpackiterable(w_state, 4) self.w_value = space.interp_w(W_UnicodeObject, w_initval) self.buf = None self.builder = None self.state = READING self._init_newline(space, w_readnl) pos = space.getindex_w(w_pos, space.w_TypeError) if pos < 0: raise oefmt(space.w_ValueError, "position value cannot be negative") self.pos = pos if not space.is_w(w_dict, space.w_None): if not space.isinstance_w(w_dict, space.w_dict): raise oefmt( space.w_TypeError, "fourth item of state should be a dict, got a %T", w_dict) # Alternatively, we could replace the internal dictionary # completely. However, it seems more practical to just update it. space.call_method(self.getdict(space), "update", w_dict) def _check_closed(self, space, message=None): if self.state == CLOSED: if message is None: message = "I/O operation on closed file" raise OperationError(space.w_ValueError, space.newtext(message)) @signature( types.self(), types.any(), types.any(), returns=types.instance(W_UnicodeObject)) def _decode_string(self, space, w_obj): if not space.isinstance_w(w_obj, space.w_unicode): raise oefmt(space.w_TypeError, "unicode argument expected, got '%T'", w_obj) self._check_closed(space) if self.w_decoder is not None: w_decoded = space.call_method( self.w_decoder, "decode", w_obj, space.w_True) else: w_decoded = w_obj writenl = self.writenl if writenl is not None: w_decoded = space.call_method( w_decoded, "replace", space.newtext("\n"), space.newutf8(writenl, codepoints_in_utf8(writenl)), ) w_decoded = space.interp_w(W_UnicodeObject, w_decoded) return w_decoded def write_w(self, space, w_obj): w_decoded = self._decode_string(space, w_obj) string, length = space.utf8_len_w(w_decoded) orig_size = space.len_w(w_obj) if not string: return space.newint(orig_size) if self.state == READING: if self.pos == self.get_length(): # switch to ACCUMULATING self.builder = Utf8StringBuilder() self.builder.append_utf8(self.w_value._utf8, self.w_value._len()) self.w_value = None self.state = ACCUMULATING else: # switch to RWBUFFER s, l = space.utf8_len_w(self.w_value) self.buf = UnicodeIO(s, l) self.w_value = None self.state = RWBUFFER elif self.state == ACCUMULATING and self.pos != self.get_length(): s = self.builder.build() self.buf = UnicodeIO(s, self.builder.getlength()) self.builder = None self.state = RWBUFFER if self.state == ACCUMULATING: written = length self.builder.append_utf8(string, length) else: assert self.state == RWBUFFER written = self.buf.write(string, length, self.pos) self.pos += written return space.newint(orig_size) def _realize(self, space): """Switch from ACCUMULATING to READING""" s = self.builder.build() length = self.builder.getlength() self.w_value = space.newutf8(s, length) self.builder = None self.state = READING def read_w(self, space, w_size=None): self._check_closed(space) size = convert_size(space, w_size) if self.pos >= self.get_length(): return W_UnicodeObject.EMPTY if self.state == ACCUMULATING: self._realize(space) if self.state == READING: length = self.w_value._len() end = _find_end(self.pos, size, length) if self.pos > end: return space.newutf8('', 0) if self.pos == 0 and end == length: self.pos = end return self.w_value w_res = self.w_value._unicode_sliced(space, self.pos, end) self.pos = end return w_res assert self.state == RWBUFFER result_builder = self.buf.read(self.pos, size) self.pos += result_builder.getlength() return W_UnicodeObject.from_utf8builder(result_builder) def readline_w(self, space, w_limit=None): self._check_closed(space) limit = convert_size(space, w_limit) if self.pos >= self.get_length(): return W_UnicodeObject.EMPTY if self.state == ACCUMULATING: self._realize(space) if self.state == READING: length = self.w_value._len() end = _find_end(self.pos, limit, length) if self.readuniversal: start = self.pos start_offset = self.w_value._index_to_byte(start) it = Utf8StringIterator(self.w_value._utf8) it._pos = start_offset for ch in it: if self.pos >= end: break if ch == ord('\n'): self.pos += 1 break elif ch == ord('\r'): self.pos += 1 if self.pos >= end: break if it.next() == ord('\n'): self.pos += 1 break else: # `it` has gone one char too far, but we don't care break self.pos += 1 w_res = self.w_value._unicode_sliced(space, start, self.pos) return w_res else: if self.readtranslate: # Newlines are already translated, only search for \n newline = '\n' else: newline = self.readnl start = self.pos start_offset = self.w_value._index_to_byte(start) it = Utf8StringIterator(self.w_value._utf8) it._pos = start_offset for ch in it: if self.pos >= end: break self.pos += 1 if ch == ord(newline[0]): if len(newline) == 1 or self.pos >= end: break else: ch = codepoint_at_pos(self.w_value._utf8, it.get_pos()) if ch == ord(newline[1]): self.pos += 1 break else: continue w_res = self.w_value._unicode_sliced(space, start, self.pos) return w_res if self.readuniversal: result_builder = self.buf.readline_universal(self.pos, limit) else: if self.readtranslate: # Newlines are already translated, only search for \n newline = '\n' else: newline = self.readnl result_builder = self.buf.readline(newline, self.pos, limit) self.pos += result_builder.getlength() return W_UnicodeObject.from_utf8builder(result_builder) @unwrap_spec(pos=int, mode=int) def seek_w(self, space, pos, mode=0): self._check_closed(space) if not 0 <= mode <= 2: raise oefmt(space.w_ValueError, "Invalid whence (%d, should be 0, 1 or 2)", mode) elif mode == 0 and pos < 0: raise oefmt(space.w_ValueError, "Negative seek position %d", pos) elif mode != 0 and pos != 0: raise oefmt(space.w_IOError, "Can't do nonzero cur-relative seeks") # XXX: this makes almost no sense, but its how CPython does it. if mode == 1: pos = self.pos elif mode == 2: pos = self.get_length() assert pos >= 0 self.pos = pos return space.newint(pos) def truncate_w(self, space, w_size=None): self._check_closed(space) if space.is_none(w_size): size = self.pos else: size = space.int_w(w_size) if size < 0: raise oefmt(space.w_ValueError, "Negative size value %d", size) if self.state == ACCUMULATING: self._realize(space) if self.state == READING: if size < self.w_value._len(): self.w_value = self.w_value._unicode_sliced(space, 0, size) else: self.buf.truncate(size) return space.newint(size) def getvalue_w(self, space): self._check_closed(space) if self.state == ACCUMULATING: self._realize(space) if self.state == READING: return self.w_value builder = self.buf.getvalue() return W_UnicodeObject.from_utf8builder(builder) def readable_w(self, space): self._check_closed(space) return space.w_True def writable_w(self, space): self._check_closed(space) return space.w_True def seekable_w(self, space): self._check_closed(space) return space.w_True def close_w(self, space): self.buf = None self.state = CLOSED def needs_finalizer(self): # 'self.buf = None' is not necessary when the object goes away return type(self) is not W_StringIO def closed_get_w(self, space): return space.newbool(self.state == CLOSED) def line_buffering_get_w(self, space): return space.w_False def newlines_get_w(self, space): if self.w_decoder is None: return space.w_None return space.getattr(self.w_decoder, space.newtext("newlines")) W_StringIO.typedef = TypeDef( '_io.StringIO', W_TextIOBase.typedef, __new__ = generic_new_descr(W_StringIO), __init__ = interp2app(W_StringIO.descr_init), __getstate__ = interp2app(W_StringIO.descr_getstate), __setstate__ = interp2app(W_StringIO.descr_setstate), write = interp2app(W_StringIO.write_w), read = interp2app(W_StringIO.read_w), readline = interp2app(W_StringIO.readline_w), seek = interp2app(W_StringIO.seek_w), truncate = interp2app(W_StringIO.truncate_w), getvalue = interp2app(W_StringIO.getvalue_w), readable = interp2app(W_StringIO.readable_w), writable = interp2app(W_StringIO.writable_w), seekable = interp2app(W_StringIO.seekable_w), close = interp2app(W_StringIO.close_w), closed = GetSetProperty(W_StringIO.closed_get_w), line_buffering = GetSetProperty(W_StringIO.line_buffering_get_w), newlines = GetSetProperty(W_StringIO.newlines_get_w), )