Implement BytesIO.peek()

Author: marcelm

Lib/_pyio.py

@@ -996,6 +996,14 @@ def tell(self):
             raise ValueError("tell on closed file")
         return self._pos
 
+    def peek(self, size=-1):
+        pos = self.tell()
+        if size == 0:
+            size = -1
+        b = self.read(size)
+        self.seek(pos)
+        return b
+
     def truncate(self, pos=None):
         if self.closed:
             raise ValueError("truncate on closed file")

Lib/test/test_io/test_memoryio.py

@@ -566,6 +566,23 @@ def test_issue141311(self):
         buf = bytearray(2)
         self.assertEqual(0, memio.readinto(buf))
 
+    def test_peek(self):
+        buf = self.buftype("1234567890")
+        memio = self.ioclass(buf)
+
+        self.assertEqual(memio.peek(1), buf[:1])
+        self.assertEqual(memio.peek(1), buf[:1])
+        self.assertEqual(memio.peek(), buf)
+        self.assertEqual(memio.peek(0), buf)
+        memio.read(1)
+        self.assertEqual(memio.peek(1), buf[1:2])
+        self.assertEqual(memio.peek(), buf[1:])
+        self.assertEqual(memio.peek(42), buf[1:])
+        memio.read()
+        self.assertEqual(memio.peek(1), self.EOF)
+        memio.close()
+        self.assertRaises(ValueError, memio.peek)
+
     def test_unicode(self):
         memio = self.ioclass()
 

Modules/_io/bytesio.c

@@ -499,6 +499,42 @@ _io_BytesIO_read1_impl(bytesio *self, Py_ssize_t size)
     return _io_BytesIO_read_impl(self, size);
 }
 
+
+/*[clinic input]
+_io.BytesIO.peek
+    size: Py_ssize_t(accept={int, NoneType}) = -1
+    /
+
+Return bytes from the stream without advancing the position.
+
+Return an empty bytes object at EOF.
+[clinic start generated code]*/
+
+static PyObject *
+_io_BytesIO_peek_impl(bytesio *self, Py_ssize_t size)
+/*[clinic end generated code: output=fa4d8ce28b35db9b input=afc80e71b37e7c59]*/
+{
+    Py_ssize_t n;
+    const char *output;
+
+    CHECK_CLOSED(self);
+
+    /* adjust invalid sizes */
+    n = self->string_size - self->pos;
+    if (size < 1 || size > n) {
+        size = n;
+        if (size < 0)
+            size = 0;
+    }
+
+    assert(self->buf != NULL);
+    assert(size <= self->string_size);
+    output = PyBytes_AS_STRING(self->buf) + self->pos;
+    return PyBytes_FromStringAndSize(output, size);
+}
+
+
+
 /*[clinic input]
 @critical_section
 _io.BytesIO.readline
@@ -1135,6 +1171,7 @@ static struct PyMethodDef bytesio_methods[] = {
     _IO_BYTESIO_READLINE_METHODDEF
     _IO_BYTESIO_READLINES_METHODDEF
     _IO_BYTESIO_READ_METHODDEF
+    _IO_BYTESIO_PEEK_METHODDEF
     _IO_BYTESIO_GETBUFFER_METHODDEF
     _IO_BYTESIO_GETVALUE_METHODDEF
     _IO_BYTESIO_SEEK_METHODDEF