[mypyc] Add minimal, experimental librt.base64 module

mypy/typeshed/stubs/librt/librt/base64.pyi

@@ -0,0 +1 @@
+def b64encode(s: bytes) -> bytes: ...

mypyc/build.py

@@ -42,7 +42,7 @@
 from mypyc.namegen import exported_name
 from mypyc.options import CompilerOptions
 
-LIBRT_MODULES = [("librt.internal", "librt_internal.c")]
+LIBRT_MODULES = [("librt.internal", "librt_internal.c"), ("librt.base64", "librt_base64.c")]
 
 try:
     # Import setuptools so that it monkey-patch overrides distutils
@@ -495,7 +495,9 @@ def mypycify(
     group_name: str | None = None,
     log_trace: bool = False,
     depends_on_librt_internal: bool = False,
+    depends_on_librt_base64: bool = False,
     install_librt: bool = False,
+    experimental_features: bool = False,
 ) -> list[Extension]:
     """Main entry point to building using mypyc.
 
@@ -551,6 +553,9 @@ def mypycify(
                        those are build and published on PyPI separately, but during
                        tests, we want to use their development versions (i.e. from
                        current commit).
+        experimental_features: Enable experimental features (install_librt=True is
+                               also needed if using experimental librt features). These
+                               have no backward compatibility guarantees!
     """
 
     # Figure out our configuration
@@ -565,6 +570,8 @@ def mypycify(
         group_name=group_name,
         log_trace=log_trace,
         depends_on_librt_internal=depends_on_librt_internal,
+        depends_on_librt_base64=depends_on_librt_base64,
+        experimental_features=experimental_features,
     )
 
     # Generate all the actual important C code
@@ -607,6 +614,8 @@ def mypycify(
         ]
         if log_trace:
             cflags.append("-DMYPYC_LOG_TRACE")
+        if experimental_features:
+            cflags.append("-DMYPYC_EXPERIMENTAL")
     elif compiler.compiler_type == "msvc":
         # msvc doesn't have levels, '/O2' is full and '/Od' is disable
         if opt_level == "0":
@@ -633,6 +642,8 @@ def mypycify(
             cflags += ["/GL-", "/wd9025"]  # warning about overriding /GL
         if log_trace:
             cflags.append("/DMYPYC_LOG_TRACE")
+        if experimental_features:
+            cflags.append("/DMYPYC_EXPERIMENTAL")
 
     # If configured to (defaults to yes in multi-file mode), copy the
     # runtime library in. Otherwise it just gets #included to save on

mypyc/codegen/emitmodule.py

@@ -604,6 +604,8 @@ def generate_c_for_modules(self) -> list[tuple[str, str]]:
         ext_declarations.emit_line("#include <CPy.h>")
         if self.compiler_options.depends_on_librt_internal:
             ext_declarations.emit_line("#include <librt_internal.h>")
+        if self.compiler_options.depends_on_librt_base64:
+            ext_declarations.emit_line("#include <librt_base64.h>")
 
         declarations = Emitter(self.context)
         declarations.emit_line(f"#ifndef MYPYC_LIBRT_INTERNAL{self.group_suffix}_H")
@@ -1034,6 +1036,10 @@ def emit_module_exec_func(
             emitter.emit_line("if (import_librt_internal() < 0) {")
             emitter.emit_line("return -1;")
             emitter.emit_line("}")
+        if self.compiler_options.depends_on_librt_base64:
+            emitter.emit_line("if (import_librt_base64() < 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/ops.py

@@ -707,6 +707,7 @@ def __init__(
         extra_int_constants: list[tuple[int, RType]],
         priority: int,
         is_pure: bool,
+        experimental: bool,
     ) -> None:
         # Each primitive much have a distinct name, but otherwise they are arbitrary.
         self.name: Final = name
@@ -729,6 +730,9 @@ def __init__(
         self.is_pure: Final = is_pure
         if is_pure:
             assert error_kind == ERR_NEVER
+        # Experimental primitives are not used unless mypyc experimental features are
+        # explicitly enabled
+        self.experimental = experimental
 
     def __repr__(self) -> str:
         return f"<PrimitiveDescription {self.name!r}: {self.arg_types}>"

mypyc/irbuild/ll_builder.py

@@ -2212,6 +2212,8 @@ def matching_primitive_op(
         for desc in candidates:
             if len(desc.arg_types) != len(args):
                 continue
+            if desc.experimental and not self.options.experimental_features:
+                continue
             if all(
                 # formal is not None and # TODO
                 is_subtype(actual.type, formal)

mypyc/lib-rt/librt_base64.c

@@ -0,0 +1,148 @@
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+#include "librt_base64.h"
+#include "pythoncapi_compat.h"
+
+#ifdef MYPYC_EXPERIMENTAL
+
+// b64encode_internal below is adapted from the CPython 3.14.0 binascii module
+
+static const unsigned char table_b2a_base64[] =
+"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
+
+#define BASE64_PAD '='
+
+/* Max binary chunk size; limited only by available memory */
+#define BASE64_MAXBIN ((PY_SSIZE_T_MAX - 3) / 2)
+
+static PyObject *
+b64encode_internal(PyObject *obj) {
+    unsigned char *ascii_data;
+    const unsigned char *bin_data;
+    int leftbits = 0;
+    unsigned char this_ch;
+    unsigned int leftchar = 0;
+    Py_ssize_t bin_len, out_len;
+    PyBytesWriter *writer;
+    int newline = 0; // TODO
+
+    if (!PyBytes_Check(obj)) {
+        PyErr_SetString(PyExc_TypeError, "base64() expects a bytes object");
+        return NULL;
+    }
+
+    bin_data = (const unsigned char *)PyBytes_AS_STRING(obj);
+    bin_len = PyBytes_GET_SIZE(obj);
+
+    assert(bin_len >= 0);
+
+    if ( bin_len > BASE64_MAXBIN ) {
+        PyErr_SetString(PyExc_ValueError, "Too much data for base64 line");
+        return NULL;
+    }
+
+    /* We're lazy and allocate too much (fixed up later).
+       "+2" leaves room for up to two pad characters.
+       Note that 'b' gets encoded as 'Yg==\n' (1 in, 5 out). */
+    out_len = bin_len*2 + 2;
+    if (newline)
+        out_len++;
+    writer = PyBytesWriter_Create(out_len);
+    ascii_data = PyBytesWriter_GetData(writer);
+    if (writer == NULL)
+        return NULL;
+
+    for( ; bin_len > 0 ; bin_len--, bin_data++ ) {
+        /* Shift the data into our buffer */
+        leftchar = (leftchar << 8) | *bin_data;
+        leftbits += 8;
+
+        /* See if there are 6-bit groups ready */
+        while ( leftbits >= 6 ) {
+            this_ch = (leftchar >> (leftbits-6)) & 0x3f;
+            leftbits -= 6;
+            *ascii_data++ = table_b2a_base64[this_ch];
+        }
+    }
+    if ( leftbits == 2 ) {
+        *ascii_data++ = table_b2a_base64[(leftchar&3) << 4];
+        *ascii_data++ = BASE64_PAD;
+        *ascii_data++ = BASE64_PAD;
+    } else if ( leftbits == 4 ) {
+        *ascii_data++ = table_b2a_base64[(leftchar&0xf) << 2];
+        *ascii_data++ = BASE64_PAD;
+    }
+    if (newline)
+        *ascii_data++ = '\n';       /* Append a courtesy newline */
+
+    return PyBytesWriter_FinishWithSize(writer, ascii_data - (unsigned char *)PyBytesWriter_GetData(writer));
+}
+
+static PyObject*
+b64encode(PyObject *self, PyObject *const *args, size_t nargs) {
+    if (nargs != 1) {
+        PyErr_SetString(PyExc_TypeError, "b64encode() takes exactly one argument");
+        return 0;
+    }
+    return b64encode_internal(args[0]);
+}
+
+#endif
+
+static PyMethodDef librt_base64_module_methods[] = {
+#ifdef MYPYC_EXPERIMENTAL
+    {"b64encode", (PyCFunction)b64encode, METH_FASTCALL, PyDoc_STR("Encode bytes-like object using Base64.")},
+#endif
+    {NULL, NULL, 0, NULL}
+};
+
+static int
+base64_abi_version(void) {
+    return 0;
+}
+
+static int
+base64_api_version(void) {
+    return 0;
+}
+
+static int
+librt_base64_module_exec(PyObject *m)
+{
+#ifdef MYPYC_EXPERIMENTAL
+    // Export mypy internal C API, be careful with the order!
+    static void *base64_api[LIBRT_BASE64_API_LEN] = {
+        (void *)base64_abi_version,
+        (void *)base64_api_version,
+        (void *)b64encode_internal,
+    };
+    PyObject *c_api_object = PyCapsule_New((void *)base64_api, "librt.base64._C_API", NULL);
+    if (PyModule_Add(m, "_C_API", c_api_object) < 0) {
+        return -1;
+    }
+#endif
+    return 0;
+}
+
+static PyModuleDef_Slot librt_base64_module_slots[] = {
+    {Py_mod_exec, librt_base64_module_exec},
+#ifdef Py_MOD_GIL_NOT_USED
+    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
+#endif
+    {0, NULL}
+};
+
+static PyModuleDef librt_base64_module = {
+    .m_base = PyModuleDef_HEAD_INIT,
+    .m_name = "base64",
+    .m_doc = "base64 encoding and decoding optimized for mypyc",
+    .m_size = 0,
+    .m_methods = librt_base64_module_methods,
+    .m_slots = librt_base64_module_slots,
+};
+
+PyMODINIT_FUNC
+PyInit_base64(void)
+{
+    return PyModuleDef_Init(&librt_base64_module);
+}

mypyc/lib-rt/librt_base64.h

@@ -0,0 +1,60 @@
+#ifndef LIBRT_BASE64_H
+#define LIBRT_BASE64_H
+
+#ifndef MYPYC_EXPERIMENTAL
+
+static int
+import_librt_base64(void)
+{
+    // All librt.base64 features are experimental for now, so don't set up the API here
+    return 0;
+}
+
+#else  // MYPYC_EXPERIMENTAL
+
+#define LIBRT_BASE64_ABI_VERSION 0
+#define LIBRT_BASE64_API_VERSION 0
+#define LIBRT_BASE64_API_LEN 3
+
+static void *LibRTBase64_API[LIBRT_BASE64_API_LEN];
+
+#define LibRTBase64_ABIVersion (*(int (*)(void)) LibRTBase64_API[0])
+#define LibRTBase64_APIVersion (*(int (*)(void)) LibRTBase64_API[1])
+#define LibRTBase64_b64encode_internal (*(PyObject* (*)(PyObject *source)) LibRTBase64_API[2])
+
+static int
+import_librt_base64(void)
+{
+    PyObject *mod = PyImport_ImportModule("librt.base64");
+    if (mod == NULL)
+        return -1;
+    Py_DECREF(mod);  // we import just for the side effect of making the below work.
+    void *capsule = PyCapsule_Import("librt.base64._C_API", 0);
+    if (capsule == NULL)
+        return -1;
+    memcpy(LibRTBase64_API, capsule, sizeof(LibRTBase64_API));
+    if (LibRTBase64_ABIVersion() != LIBRT_BASE64_ABI_VERSION) {
+        char err[128];
+        snprintf(err, sizeof(err), "ABI version conflict for librt.base64, expected %d, found %d",
+            LIBRT_BASE64_ABI_VERSION,
+            LibRTBase64_ABIVersion()
+        );
+        PyErr_SetString(PyExc_ValueError, err);
+        return -1;
+    }
+    if (LibRTBase64_APIVersion() < LIBRT_BASE64_API_VERSION) {
+        char err[128];
+        snprintf(err, sizeof(err),
+                 "API version conflict for librt.base64, expected %d or newer, found %d (hint: upgrade librt)",
+            LIBRT_BASE64_API_VERSION,
+            LibRTBase64_APIVersion()
+        );
+        PyErr_SetString(PyExc_ValueError, err);
+        return -1;
+    }
+    return 0;
+}
+
+#endif  // MYPYC_EXPERIMENTAL
+
+#endif  // LIBRT_BASE64_H

mypyc/lib-rt/setup.py

@@ -96,6 +96,9 @@ def run(self) -> None:
                 ],
                 include_dirs=["."],
                 extra_compile_args=cflags,
-            )
+            ),
+            Extension(
+                "librt.base64", ["librt_base64.c"], include_dirs=["."], extra_compile_args=cflags
+            ),
         ]
     )

mypyc/options.py

@@ -18,6 +18,8 @@ def __init__(
         group_name: str | None = None,
         log_trace: bool = False,
         depends_on_librt_internal: bool = False,
+        depends_on_librt_base64: bool = False,
+        experimental_features: bool = False,
     ) -> None:
         self.strip_asserts = strip_asserts
         self.multi_file = multi_file
@@ -55,3 +57,8 @@ def __init__(
         # only for mypy itself, third-party code compiled with mypyc should not use
         # librt.internal.
         self.depends_on_librt_internal = depends_on_librt_internal
+        self.depends_on_librt_base64 = depends_on_librt_base64
+        # Some experimental features are only available when building librt in
+        # experimental mode (e.g. use _experimental suffix in librt run test).
+        # These can't be used with a librt wheel installed from PyPI.
+        self.experimental_features = experimental_features

mypyc/primitives/misc_ops.py

@@ -465,3 +465,12 @@
     c_function_name="cache_version_internal",
     error_kind=ERR_NEVER,
 )
+
+function_op(
+    name="librt.base64.b64encode",
+    arg_types=[bytes_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name="LibRTBase64_b64encode_internal",
+    error_kind=ERR_MAGIC,
+    experimental=True,
+)