Add PyBytesWriter_GetRemaining()

Author: vstinner

Doc/c-api/bytes.rst

@@ -284,6 +284,14 @@ PyBytesWriter
    On success, return non-``NULL`` buffer.
    On error, set an exception and return ``NULL``.
 
+.. c:function:: Py_ssize_t PyBytesWriter_GetRemaining(PyBytesWriter *writer, void *buf)
+
+   Get the number of remaining bytes to write.
+
+   It is the difference between total allocated bytes (bytes allocated by
+   :c:func:`PyBytesWriter_Create` and :c:func:`PyBytesWriter_Extend`) and the
+   current position in the buffer.
+
 .. c:function:: Py_ssize_t PyBytesWriter_GetAllocated(PyBytesWriter *writer)
 
    Get the number of allocated bytes.

Include/cpython/bytesobject.h

@@ -55,6 +55,9 @@ PyAPI_FUNC(PyObject*) PyBytesWriter_Finish(
     PyBytesWriter *writer,
     void *buf);
 
+PyAPI_FUNC(Py_ssize_t) PyBytesWriter_GetRemaining(
+    PyBytesWriter *writer,
+    void *buf);
 PyAPI_FUNC(Py_ssize_t) PyBytesWriter_GetAllocated(
     PyBytesWriter *writer);
 PyAPI_FUNC(void*) PyBytesWriter_Extend(

Lib/test/test_capi/test_bytes.py

@@ -297,6 +297,7 @@ def create_writer(self, alloc):
     def test_empty(self):
         # Test PyBytesWriter_Create()
         writer = self.create_writer(0)
+        self.assertEqual(writer.get_remaining(), 0)
         self.assertEqual(writer.finish(), b'')
 
     def test_write_bytes(self):
@@ -317,29 +318,34 @@ def test_extend(self):
         writer = self.create_writer(0)
         writer.extend(20, b'number=123456')
         writer.extend(0, b'')
+        self.assertEqual(writer.get_remaining(), 7)
         self.assertEqual(writer.finish(), b'number=123456')
 
         writer = self.create_writer(0)
         writer.extend(0, b'')
         writer.extend(20, b'number=123456')
+        self.assertEqual(writer.get_remaining(), 7)
         self.assertEqual(writer.finish(), b'number=123456')
 
         writer = self.create_writer(0)
         writer.extend(10, b'number=')
         writer.extend(10, b'123456')
+        self.assertEqual(writer.get_remaining(), 7)
         self.assertEqual(writer.finish(), b'number=123456')
 
         writer = self.create_writer(0)
         writer.extend(10, b'number=')
         writer.extend(0, b'')
         writer.extend(10, b'123456')
+        self.assertEqual(writer.get_remaining(), 7)
         self.assertEqual(writer.finish(), b'number=123456')
 
         writer = self.create_writer(0)
         writer.extend(10, b'number')
         writer.extend(10, b'=')
         writer.extend(10, b'123')
         writer.extend(10, b'456')
+        self.assertEqual(writer.get_remaining(), 27)
         self.assertEqual(writer.finish(), b'number=123456')
 
     def test_format(self):

Modules/_testcapi/bytes.c

@@ -196,6 +196,19 @@ writer_extend(PyObject *self_raw, PyObject *args)
 }
 
 
+static PyObject*
+writer_get_remaining(PyObject *self_raw, PyObject *Py_UNUSED(args))
+{
+    WriterObject *self = (WriterObject *)self_raw;
+    if (writer_check(self) < 0) {
+        return NULL;
+    }
+
+    Py_ssize_t size = PyBytesWriter_GetRemaining(self->writer, self->buf);
+    return PyLong_FromSsize_t(size);
+}
+
+
 static PyObject*
 writer_finish(PyObject *self_raw, PyObject *Py_UNUSED(args))
 {
@@ -214,6 +227,7 @@ static PyMethodDef writer_methods[] = {
     {"write_bytes", _PyCFunction_CAST(writer_write_bytes), METH_VARARGS},
     {"format_i", _PyCFunction_CAST(writer_format_i), METH_VARARGS},
     {"extend", _PyCFunction_CAST(writer_extend), METH_VARARGS},
+    {"get_remaining", _PyCFunction_CAST(writer_get_remaining), METH_NOARGS},
     {"finish", _PyCFunction_CAST(writer_finish), METH_NOARGS},
     {NULL,              NULL}           /* sentinel */
 };

Objects/bytesobject.c

@@ -3627,6 +3627,14 @@ PyBytesWriter_Finish(PyBytesWriter *writer, void *buf)
 }
 
 
+Py_ssize_t
+PyBytesWriter_GetRemaining(PyBytesWriter *writer, void *buf)
+{
+    Py_ssize_t pos = (char*)buf - byteswriter_start(writer);
+    return writer->size - pos;
+}
+
+
 Py_ssize_t
 PyBytesWriter_GetAllocated(PyBytesWriter *writer)
 {