From c755c4a666ab3b327490d3a585ec4935f54e0fcb Mon Sep 17 00:00:00 2001
From: Marcel Martin <marcel.martin@scilifelab.se>
Date: Sat, 22 Jan 2022 22:21:41 +0100
Subject: [PATCH] Implement BytesIO.peek()

---
 Lib/_pyio.py                   |  8 +++++++
 Lib/test/test_memoryio.py      | 17 +++++++++++++++
 Modules/_io/bytesio.c          | 37 +++++++++++++++++++++++++++++++++
 Modules/_io/clinic/bytesio.c.h | 38 +++++++++++++++++++++++++++++++++-
 4 files changed, 99 insertions(+), 1 deletion(-)

diff --git a/Lib/_pyio.py b/Lib/_pyio.py
index 12510784c8b9269..54fbee409320268 100644
--- a/Lib/_pyio.py
+++ b/Lib/_pyio.py
@@ -993,6 +993,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")
diff --git a/Lib/test/test_memoryio.py b/Lib/test/test_memoryio.py
index cd2faba1791c772..f77a99f1bbe14dc 100644
--- a/Lib/test/test_memoryio.py
+++ b/Lib/test/test_memoryio.py
@@ -517,6 +517,23 @@ def test_relative_seek(self):
         memio.seek(1, 1)
         self.assertEqual(memio.read(), buf[1:])
 
+    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()
 
diff --git a/Modules/_io/bytesio.c b/Modules/_io/bytesio.c
index 930ef7e29dbcf67..386ec621207eac4 100644
--- a/Modules/_io/bytesio.c
+++ b/Modules/_io/bytesio.c
@@ -459,6 +459,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]
 _io.BytesIO.readline
     size: Py_ssize_t(accept={int, NoneType}) = -1
@@ -1014,6 +1050,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
diff --git a/Modules/_io/clinic/bytesio.c.h b/Modules/_io/clinic/bytesio.c.h
index 049d3473110f7b6..15a0235441d6d13 100644
--- a/Modules/_io/clinic/bytesio.c.h
+++ b/Modules/_io/clinic/bytesio.c.h
@@ -222,6 +222,42 @@ _io_BytesIO_read1(bytesio *self, PyObject *const *args, Py_ssize_t nargs)
     return return_value;
 }
 
+PyDoc_STRVAR(_io_BytesIO_peek__doc__,
+"peek($self, size=-1, /)\n"
+"--\n"
+"\n"
+"Return bytes from the stream without advancing the position.\n"
+"\n"
+"Return an empty bytes object at EOF.");
+
+#define _IO_BYTESIO_PEEK_METHODDEF    \
+    {"peek", _PyCFunction_CAST(_io_BytesIO_peek), METH_FASTCALL, _io_BytesIO_peek__doc__},
+
+static PyObject *
+_io_BytesIO_peek_impl(bytesio *self, Py_ssize_t size);
+
+static PyObject *
+_io_BytesIO_peek(bytesio *self, PyObject *const *args, Py_ssize_t nargs)
+{
+    PyObject *return_value = NULL;
+    Py_ssize_t size = -1;
+
+    if (!_PyArg_CheckPositional("peek", nargs, 0, 1)) {
+        goto exit;
+    }
+    if (nargs < 1) {
+        goto skip_optional;
+    }
+    if (!_Py_convert_optional_to_ssize_t(args[0], &size)) {
+        goto exit;
+    }
+skip_optional:
+    return_value = _io_BytesIO_peek_impl(self, size);
+
+exit:
+    return return_value;
+}
+
 PyDoc_STRVAR(_io_BytesIO_readline__doc__,
 "readline($self, size=-1, /)\n"
 "--\n"
@@ -505,4 +541,4 @@ _io_BytesIO___init__(PyObject *self, PyObject *args, PyObject *kwargs)
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=93d9700a6cf395b8 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=8bd1b169f6439bb8 input=a9049054013a1b77]*/