[mypyc] Add minimal BytesWriter class librt (#20297)

The `librt.strings.BytesWriter` class allow building a `bytes` object
from `bytes` or `u8` components.

Simple example use case:
```
from librt.strings import BytesWriter

w = BytesWriter()
w.write(b'foo')
w.append(ord('!'))
print(w.getvalue())  # b'foo!'
```

It can be used at an alternative to `bytearray` or `BytesIO` in various
use cases. The main benefit over stdlib classes is that we can provide
(more) efficient mypyc primitives for various operations. The goal is to
make it more efficient than either `bytearray` or `BytesIO` in many use
cases, eventually.

The implementation is based on the implementation of
`librt.internal.WriteBuffer`, but it is somewhat different, since it
will support also direct indexed item access (a bit like `bytearray`),
so just appending to the end is not the only important use case. Direct
item access is not implemented yet, but I'll add more features in
follow-up PRs.

This minimal implementation hasn't been optimized yet, and the API may
not be final. It's experimental and only compiled if experimental
features are explicitly enabled, so we can still make arbitrary API and
ABI changes.

Author: JukkaL

mypy/typeshed/stubs/librt/librt/strings.pyi

@@ -0,0 +1,7 @@
+from typing import final
+
+@final
+class BytesWriter:
+    def append(self, /, x: int) -> None: ...
+    def write(self, /, b: bytes) -> None: ...
+    def getvalue(self) -> bytes: ...

mypyc/build.py

@@ -52,6 +52,7 @@ class ModDesc(NamedTuple):
 
 LIBRT_MODULES = [
     ModDesc("librt.internal", ["librt_internal.c"], [], []),
+    ModDesc("librt.strings", ["librt_strings.c"], [], []),
     ModDesc(
         "librt.base64",
         [

mypyc/codegen/emitmodule.py

@@ -611,6 +611,8 @@ def generate_c_for_modules(self) -> list[tuple[str, str]]:
             ext_declarations.emit_line("#include <librt_internal.h>")
         if any("librt.base64" in mod.capsules for mod in self.modules.values()):
             ext_declarations.emit_line("#include <librt_base64.h>")
+        if any("librt.strings" in mod.capsules for mod in self.modules.values()):
+            ext_declarations.emit_line("#include <librt_strings.h>")
 
         declarations = Emitter(self.context)
         declarations.emit_line(f"#ifndef MYPYC_LIBRT_INTERNAL{self.group_suffix}_H")
@@ -1045,6 +1047,10 @@ def emit_module_exec_func(
             emitter.emit_line("if (import_librt_base64() < 0) {")
             emitter.emit_line("return -1;")
             emitter.emit_line("}")
+        if "librt.strings" in module.capsules:
+            emitter.emit_line("if (import_librt_strings() < 0) {")
+            emitter.emit_line("return -1;")
+            emitter.emit_line("}")
         emitter.emit_line("PyObject* modname = NULL;")
         if self.multi_phase_init:
             emitter.emit_line(f"{module_static} = module;")

mypyc/ir/rtypes.py

@@ -514,7 +514,11 @@ def __hash__(self) -> int:
 
 KNOWN_NATIVE_TYPES: Final = {
     name: RPrimitive(name, is_unboxed=False, is_refcounted=True)
-    for name in ["librt.internal.WriteBuffer", "librt.internal.ReadBuffer"]
+    for name in [
+        "librt.internal.WriteBuffer",
+        "librt.internal.ReadBuffer",
+        "librt.strings.BytesWriter",
+    ]
 }
 
 

mypyc/lib-rt/librt_strings.c

@@ -0,0 +1,344 @@
+#include "pythoncapi_compat.h"
+
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+#include <stdint.h>
+#include "CPy.h"
+#include "librt_strings.h"
+
+#define CPY_BOOL_ERROR 2
+#define CPY_NONE_ERROR 2
+#define CPY_NONE 1
+
+//
+// BytesWriter
+//
+
+// Length of the default buffer embedded directly in a BytesWriter object
+#define WRITER_EMBEDDED_BUF_LEN 512
+
+typedef struct {
+    PyObject_HEAD
+    char *buf;  // Beginning of the buffer
+    Py_ssize_t len;  // Current length (number of bytes written)
+    Py_ssize_t capacity;  // Total capacity of the buffer
+    char data[WRITER_EMBEDDED_BUF_LEN];  // Default buffer
+} BytesWriterObject;
+
+#define _WRITE(data, type, v) \
+    do { \
+       *(type *)(((BytesWriterObject *)data)->buf + ((BytesWriterObject *)data)->len) = v; \
+       ((BytesWriterObject *)data)->len += sizeof(type); \
+    } while (0)
+
+static PyTypeObject BytesWriterType;
+
+static bool
+_grow_buffer(BytesWriterObject *data, Py_ssize_t n) {
+    Py_ssize_t target = data->len + n;
+    Py_ssize_t size = data->capacity;
+    Py_ssize_t old_size = size;
+    do {
+        size *= 2;
+    } while (target >= size);
+    if (old_size == WRITER_EMBEDDED_BUF_LEN) {
+        // Move from embedded buffer to heap-allocated buffer
+        data->buf = PyMem_Malloc(size);
+        if (data->buf != NULL) {
+            memcpy(data->buf, data->data, WRITER_EMBEDDED_BUF_LEN);
+        }
+    } else {
+        data->buf = PyMem_Realloc(data->buf, size);
+    }
+    if (unlikely(data->buf == NULL)) {
+        PyErr_NoMemory();
+        return false;
+    }
+    data->capacity = size;
+    return true;
+}
+
+static inline bool
+ensure_bytes_writer_size(BytesWriterObject *data, Py_ssize_t n) {
+    if (likely(data->capacity - data->len >= n)) {
+        return true;
+    } else {
+        return _grow_buffer(data, n);
+    }
+}
+
+static inline void
+BytesWriter_init_internal(BytesWriterObject *self) {
+    self->buf = self->data;
+    self->len = 0;
+    self->capacity = WRITER_EMBEDDED_BUF_LEN;
+}
+
+static PyObject*
+BytesWriter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
+{
+    if (type != &BytesWriterType) {
+        PyErr_SetString(PyExc_TypeError, "BytesWriter cannot be subclassed");
+        return NULL;
+    }
+
+    BytesWriterObject *self = (BytesWriterObject *)type->tp_alloc(type, 0);
+    if (self != NULL) {
+        BytesWriter_init_internal(self);
+    }
+    return (PyObject *)self;
+}
+
+static PyObject *
+BytesWriter_internal(void) {
+    BytesWriterObject *self = (BytesWriterObject *)BytesWriterType.tp_alloc(&BytesWriterType, 0);
+    if (self == NULL)
+        return NULL;
+    BytesWriter_init_internal(self);
+    return (PyObject *)self;
+}
+
+static int
+BytesWriter_init(BytesWriterObject *self, PyObject *args, PyObject *kwds)
+{
+    if (!PyArg_ParseTuple(args, "")) {
+        return -1;
+    }
+
+    if (kwds != NULL && PyDict_Size(kwds) > 0) {
+        PyErr_SetString(PyExc_TypeError,
+                        "BytesWriter() takes no keyword arguments");
+        return -1;
+    }
+
+    BytesWriter_init_internal(self);
+    return 0;
+}
+
+static void
+BytesWriter_dealloc(BytesWriterObject *self)
+{
+    if (self->buf != self->data) {
+        PyMem_Free(self->buf);
+        self->buf = NULL;
+    }
+    Py_TYPE(self)->tp_free((PyObject *)self);
+}
+
+static PyObject*
+BytesWriter_getvalue_internal(PyObject *self)
+{
+    BytesWriterObject *obj = (BytesWriterObject *)self;
+    return PyBytes_FromStringAndSize(obj->buf, obj->len);
+}
+
+static PyObject*
+BytesWriter_repr(BytesWriterObject *self)
+{
+    PyObject *value = BytesWriter_getvalue_internal((PyObject *)self);
+    if (value == NULL) {
+        return NULL;
+    }
+    PyObject *value_repr = PyObject_Repr(value);
+    Py_DECREF(value);
+    if (value_repr == NULL) {
+        return NULL;
+    }
+    PyObject *result = PyUnicode_FromFormat("BytesWriter(%U)", value_repr);
+    Py_DECREF(value_repr);
+    return result;
+}
+
+static PyObject*
+BytesWriter_getvalue(BytesWriterObject *self, PyObject *Py_UNUSED(ignored))
+{
+    return PyBytes_FromStringAndSize(self->buf, self->len);
+}
+
+static PyObject* BytesWriter_append(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames);
+static PyObject* BytesWriter_write(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames);
+
+static PyMethodDef BytesWriter_methods[] = {
+    {"append", (PyCFunction) BytesWriter_append, METH_FASTCALL | METH_KEYWORDS,
+     PyDoc_STR("Append a single byte to the buffer")
+    },
+    {"write", (PyCFunction) BytesWriter_write, METH_FASTCALL | METH_KEYWORDS,
+     PyDoc_STR("Append bytes to the buffer")
+    },
+    {"getvalue", (PyCFunction) BytesWriter_getvalue, METH_NOARGS,
+     "Return the buffer content as bytes object"
+    },
+    {NULL}  /* Sentinel */
+};
+
+static PyTypeObject BytesWriterType = {
+    .ob_base = PyVarObject_HEAD_INIT(NULL, 0)
+    .tp_name = "BytesWriter",
+    .tp_doc = PyDoc_STR("Memory buffer for building bytes objects from parts"),
+    .tp_basicsize = sizeof(BytesWriterObject),
+    .tp_itemsize = 0,
+    .tp_flags = Py_TPFLAGS_DEFAULT,
+    .tp_new = BytesWriter_new,
+    .tp_init = (initproc) BytesWriter_init,
+    .tp_dealloc = (destructor) BytesWriter_dealloc,
+    .tp_methods = BytesWriter_methods,
+    .tp_repr = (reprfunc)BytesWriter_repr,
+};
+
+static inline bool
+check_bytes_writer(PyObject *data) {
+    if (unlikely(Py_TYPE(data) != &BytesWriterType)) {
+        PyErr_Format(
+            PyExc_TypeError, "data must be a BytesWriter object, got %s", Py_TYPE(data)->tp_name
+        );
+        return false;
+    }
+    return true;
+}
+
+static char
+BytesWriter_write_internal(BytesWriterObject *self, PyObject *value) {
+    const char *data;
+    Py_ssize_t size;
+    if (likely(PyBytes_Check(value))) {
+        data = PyBytes_AS_STRING(value);
+        size = PyBytes_GET_SIZE(value);
+    } else {
+        data = PyByteArray_AS_STRING(value);
+        size = PyByteArray_GET_SIZE(value);
+    }
+    // Write bytes content.
+    if (!ensure_bytes_writer_size(self, size))
+        return CPY_NONE_ERROR;
+    memcpy(self->buf + self->len, data, size);
+    self->len += size;
+    return CPY_NONE;
+}
+
+static PyObject*
+BytesWriter_write(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
+    static const char * const kwlist[] = {"value", 0};
+    static CPyArg_Parser parser = {"O:write", kwlist, 0};
+    PyObject *value;
+    if (unlikely(!CPyArg_ParseStackAndKeywordsSimple(args, nargs, kwnames, &parser, &value))) {
+        return NULL;
+    }
+    if (!check_bytes_writer(self)) {
+        return NULL;
+    }
+    if (unlikely(!PyBytes_Check(value) && !PyByteArray_Check(value))) {
+        PyErr_SetString(PyExc_TypeError, "value must be a bytes or bytearray object");
+        return NULL;
+    }
+    if (unlikely(BytesWriter_write_internal((BytesWriterObject *)self, value) == CPY_NONE_ERROR)) {
+        return NULL;
+    }
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
+static inline char
+BytesWriter_append_internal(BytesWriterObject *self, uint8_t value) {
+    if (!ensure_bytes_writer_size(self, 1))
+        return CPY_NONE_ERROR;
+    _WRITE(self, uint8_t, value);
+    return CPY_NONE;
+}
+
+static PyObject*
+BytesWriter_append(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
+    static const char * const kwlist[] = {"value", 0};
+    static CPyArg_Parser parser = {"O:append", kwlist, 0};
+    PyObject *value;
+    if (unlikely(!CPyArg_ParseStackAndKeywordsSimple(args, nargs, kwnames, &parser, &value))) {
+        return NULL;
+    }
+    if (!check_bytes_writer(self)) {
+        return NULL;
+    }
+    uint8_t unboxed = CPyLong_AsUInt8(value);
+    if (unlikely(unboxed == CPY_LL_UINT_ERROR && PyErr_Occurred())) {
+        CPy_TypeError("u8", value);
+        return NULL;
+    }
+    if (unlikely(BytesWriter_append_internal((BytesWriterObject *)self, unboxed) == CPY_NONE_ERROR)) {
+        return NULL;
+    }
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
+static PyTypeObject *
+BytesWriter_type_internal(void) {
+    return &BytesWriterType;  // Return borrowed reference
+};
+
+static PyMethodDef librt_strings_module_methods[] = {
+    {NULL, NULL, 0, NULL}
+};
+
+#ifdef MYPYC_EXPERIMENTAL
+
+static int
+strings_abi_version(void) {
+    return LIBRT_STRINGS_ABI_VERSION;
+}
+
+static int
+strings_api_version(void) {
+    return LIBRT_STRINGS_API_VERSION;
+}
+
+#endif
+
+static int
+librt_strings_module_exec(PyObject *m)
+{
+#ifdef MYPYC_EXPERIMENTAL
+    if (PyType_Ready(&BytesWriterType) < 0) {
+        return -1;
+    }
+    if (PyModule_AddObjectRef(m, "BytesWriter", (PyObject *) &BytesWriterType) < 0) {
+        return -1;
+    }
+
+    // Export mypy internal C API, be careful with the order!
+    static void *librt_strings_api[LIBRT_STRINGS_API_LEN] = {
+        (void *)strings_abi_version,
+        (void *)strings_api_version,
+        (void *)BytesWriter_internal,
+        (void *)BytesWriter_getvalue_internal,
+        (void *)BytesWriter_append_internal,
+        (void *)BytesWriter_write_internal,
+        (void *)BytesWriter_type_internal,
+    };
+    PyObject *c_api_object = PyCapsule_New((void *)librt_strings_api, "librt.strings._C_API", NULL);
+    if (PyModule_Add(m, "_C_API", c_api_object) < 0) {
+        return -1;
+    }
+#endif
+    return 0;
+}
+
+static PyModuleDef_Slot librt_strings_module_slots[] = {
+    {Py_mod_exec, librt_strings_module_exec},
+#ifdef Py_MOD_GIL_NOT_USED
+    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
+#endif
+    {0, NULL}
+};
+
+static PyModuleDef librt_strings_module = {
+    .m_base = PyModuleDef_HEAD_INIT,
+    .m_name = "strings",
+    .m_doc = "Utilities for working with str and bytes objects",
+    .m_size = 0,
+    .m_methods = librt_strings_module_methods,
+    .m_slots = librt_strings_module_slots,
+};
+
+PyMODINIT_FUNC
+PyInit_strings(void)
+{
+    return PyModuleDef_Init(&librt_strings_module);
+}