[pypy-commit] pypy py3.5: Bug-for-bug compatibility (and performance optimisation) in BufferedReader.readinto1()

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: py3.5
Changeset: r93102:feaba8e9bb0a
Date: 2017-11-20 18:34 +0000
http://bitbucket.org/pypy/pypy/changeset/feaba8e9bb0a/

Log:	Bug-for-bug compatibility (and performance optimisation) in
	BufferedReader.readinto1()

diff --git a/pypy/module/_io/interp_bufferedio.py b/pypy/module/_io/interp_bufferedio.py
--- a/pypy/module/_io/interp_bufferedio.py
+++ b/pypy/module/_io/interp_bufferedio.py
@@ -111,14 +111,15 @@
         self._unsupportedoperation(space, "detach")
 
     def readinto_w(self, space, w_buffer):
-        return self._readinto(space, w_buffer, "read")
+        return self._readinto(space, w_buffer, read_once=False)
 
     def readinto1_w(self, space, w_buffer):
-        return self._readinto(space, w_buffer, "read1")
+        return self._readinto(space, w_buffer, read_once=True)
 
-    def _readinto(self, space, w_buffer, methodname):
+    def _readinto(self, space, w_buffer, read_once):
         rwbuffer = space.writebuf_w(w_buffer)
         length = rwbuffer.getlength()
+        methodname = "read1" if read_once else "read"
         w_data = space.call_method(self, methodname, space.newint(length))
 
         if not space.isinstance_w(w_data, space.w_bytes):
@@ -882,6 +883,52 @@
         self._reader_reset_buf()
         self.state = STATE_OK
 
+    def _readinto(self, space, w_buffer, read_once):
+        rwbuffer = space.writebuf_w(w_buffer)
+        length = rwbuffer.getlength()
+        with self.lock:
+            have = self._readahead()
+            if have >= length:
+                rwbuffer.setslice(0, self.buffer[self.pos:self.pos + length])
+                return space.newint(length)
+            written = 0
+            if have > 0:
+                rwbuffer.setslice(0, self.buffer[self.pos:self.read_end])
+                written = have
+
+            while written < length:
+                if self.writable:
+                    self._flush_and_rewind_unlocked(space)
+                self._reader_reset_buf()
+                self.pos = 0
+                if written + len(self.buffer) < length:
+                    try:
+                        got = self._raw_read(space, rwbuffer, written, length - written)
+                        written += got
+                    except BlockingIOError:
+                        got = 0
+                    if got == 0:
+                        break
+                elif read_once and written:
+                    break
+                else:
+                    try:
+                        have = self._fill_buffer(space)
+                    except BlockingIOError:
+                        have = 0
+                    if have == 0:
+                        break
+                    endpos = min(have, length - written)
+                    assert endpos >= 0
+                    rwbuffer.setslice(written, self.buffer[0:endpos])
+                    written += endpos
+                    self.pos = endpos
+                if read_once:
+                    break
+            return space.newint(written)
+
+
+
 W_BufferedReader.typedef = TypeDef(
     '_io.BufferedReader', W_BufferedIOBase.typedef,
     __new__ = generic_new_descr(W_BufferedReader),
diff --git a/pypy/module/_io/test/test_bufferedio.py b/pypy/module/_io/test/test_bufferedio.py
--- a/pypy/module/_io/test/test_bufferedio.py
+++ b/pypy/module/_io/test/test_bufferedio.py
@@ -189,6 +189,31 @@
         b = bytearray(2)
         raises(ValueError, bufio.readinto, b)
 
+    def test_readinto1(self):
+        import _io
+
+        class MockIO(_io._IOBase):
+            def readable(self):
+                return True
+
+            def readinto(self, buf):
+                buf[:3] = b"abc"
+                return 3
+        bufio = _io.BufferedReader(MockIO(), buffer_size=5)
+        buf = bytearray(10)
+        bufio.read(2)
+        n = bufio.readinto1(buf)
+        assert n == 4
+        assert buf[:n] == b'cabc'
+
+        # Yes, CPython's observable behavior depends on buffer_size!
+        bufio = _io.BufferedReader(MockIO(), buffer_size=20)
+        buf = bytearray(10)
+        bufio.read(2)
+        n = bufio.readinto1(buf)
+        assert n == 1
+        assert buf[:n] == b'c'
+
     def test_seek(self):
         import _io
         raw = _io.FileIO(self.tmpfile)