[mypyc] Add minimal, experimental librt.base64 module (#20226)

The module currently only has a `b64encode` function adapted from
CPython. It's only enabled when librt is compiled with experimental
features enabled, so that we are free to iterate on this and break
backward compatibility until we are ready to declare the module as
stable.

This also adds a way to define experimental features in `librt` (and
mypyc in general, but it's currently only used for `librt`).

In follow-up PRs I'm planning to add a more efficient implementation of
`b64encode` and add more features to the module, including decoding. I'm
not planning to include every feature from the stdlib base64 module,
since many of them aren't used very widely.

Author: JukkaL

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,
+)