[mypyc] Add minimal BytesWriter class librt

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);
+}