From 148d2f71fd8b9311f1bb3104648db6c78640eaad Mon Sep 17 00:00:00 2001
From: Cody Maloney <cmaloney@theoreticalchaos.com>
Date: Wed, 22 Jan 2025 16:46:31 -0800
Subject: [PATCH] gh-129205: Add os.readinto API

Add a new OS api which will read data directly into a caller provided
writeable buffer protocol object.
---
 Modules/clinic/posixmodule.c.h | 53 +++++++++++++++++++++++++++++++++-
 Modules/posixmodule.c          | 24 +++++++++++++++
 2 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/Modules/clinic/posixmodule.c.h b/Modules/clinic/posixmodule.c.h
index 96bf21dced92f0..e0c6116ce89524 100644
--- a/Modules/clinic/posixmodule.c.h
+++ b/Modules/clinic/posixmodule.c.h
@@ -7577,6 +7577,57 @@ os_read(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
     return return_value;
 }
 
+PyDoc_STRVAR(os_readinto__doc__,
+"readinto($module, fd, buffer, /)\n"
+"--\n"
+"\n"
+"Read into a :ref:`buffer protocol <bufferobjects>` object from a file descriptor.\n"
+"\n"
+"The buffer should be mutable and accept bytes. On success, returns the number of\n"
+"bytes read. Less bytes may be read than the size of the buffer. Will retry the\n"
+"underlying system call when interrupted by a signal. For other errors, the\n"
+"system call will not be retried.");
+
+#define OS_READINTO_METHODDEF    \
+    {"readinto", _PyCFunction_CAST(os_readinto), METH_FASTCALL, os_readinto__doc__},
+
+static Py_ssize_t
+os_readinto_impl(PyObject *module, int fd, Py_buffer *buffer);
+
+static PyObject *
+os_readinto(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
+{
+    PyObject *return_value = NULL;
+    int fd;
+    Py_buffer buffer = {NULL, NULL};
+    Py_ssize_t _return_value;
+
+    if (!_PyArg_CheckPositional("readinto", nargs, 2, 2)) {
+        goto exit;
+    }
+    fd = PyLong_AsInt(args[0]);
+    if (fd == -1 && PyErr_Occurred()) {
+        goto exit;
+    }
+    if (PyObject_GetBuffer(args[1], &buffer, PyBUF_WRITABLE) < 0) {
+        _PyArg_BadArgument("readinto", "argument 2", "read-write bytes-like object", args[1]);
+        goto exit;
+    }
+    _return_value = os_readinto_impl(module, fd, &buffer);
+    if ((_return_value == -1) && PyErr_Occurred()) {
+        goto exit;
+    }
+    return_value = PyLong_FromSsize_t(_return_value);
+
+exit:
+    /* Cleanup for buffer */
+    if (buffer.obj) {
+       PyBuffer_Release(&buffer);
+    }
+
+    return return_value;
+}
+
 #if defined(HAVE_READV)
 
 PyDoc_STRVAR(os_readv__doc__,
@@ -13140,4 +13191,4 @@ os__emscripten_debugger(PyObject *module, PyObject *Py_UNUSED(ignored))
 #ifndef OS__EMSCRIPTEN_DEBUGGER_METHODDEF
     #define OS__EMSCRIPTEN_DEBUGGER_METHODDEF
 #endif /* !defined(OS__EMSCRIPTEN_DEBUGGER_METHODDEF) */
-/*[clinic end generated code: output=34cb96bd07bcef90 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=c2f04dda4ea1a399 input=a9049054013a1b77]*/
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
index fb9e55a57703fc..791f8dc73df829 100644
--- a/Modules/posixmodule.c
+++ b/Modules/posixmodule.c
@@ -11433,6 +11433,29 @@ os_read_impl(PyObject *module, int fd, Py_ssize_t length)
     return buffer;
 }
 
+/*[clinic input]
+os.readinto -> Py_ssize_t
+    fd: int
+    buffer: Py_buffer(accept={rwbuffer})
+    /
+
+Read into a :ref:`buffer protocol <bufferobjects>` object from a file descriptor.
+
+The buffer should be mutable and accept bytes. On success, returns the number of
+bytes read. Less bytes may be read than the size of the buffer. Will retry the
+underlying system call when interrupted by a signal. For other errors, the
+system call will not be retried.
+[clinic start generated code]*/
+
+static Py_ssize_t
+os_readinto_impl(PyObject *module, int fd, Py_buffer *buffer)
+/*[clinic end generated code: output=8091a3513c683a80 input=810c820f4d9b1c6b]*/
+{
+    // Cap to max read size to prevent overflow in cast to size_t for _Py_read.
+    size_t length = Py_MIN(buffer->len, _PY_READ_MAX);
+    return _Py_read(fd, buffer->buf, length);
+}
+
 #if (defined(HAVE_SENDFILE) && (defined(__FreeBSD__) || defined(__DragonFly__) \
                                 || defined(__APPLE__))) \
     || defined(HAVE_READV) || defined(HAVE_PREADV) || defined (HAVE_PREADV2) \
@@ -16973,6 +16996,7 @@ static PyMethodDef posix_methods[] = {
     OS_LOCKF_METHODDEF
     OS_LSEEK_METHODDEF
     OS_READ_METHODDEF
+    OS_READINTO_METHODDEF
     OS_READV_METHODDEF
     OS_PREAD_METHODDEF
     OS_PREADV_METHODDEF