Ziggy Pydust - Python, Poetry and Zig

build.py

from setuptools_zig import ZigExtension, zig_build_ext


def build(setup_kwargs):
    setup_kwargs["ext_modules"] = [ZigExtension("mypackage._mymodule", build_zig="path/to/build.zig")]
    setup_kwargs["cmdclass"] = {"build_ext": zig_build_ext}

build.zig

const std = @import("std");

pub fn build(b: *std.Build) void {
    const lib = b.addSharedLibrary(.{
        .name = "mymodule",
        .root_source_file = .{ .path = "src/main.zig" },
        .target = b.standardTargetOptions(.{}),
        .optimize = b.standardOptimizeOption(.{}),
    });
    // Allow undefined symbols when linking, allowing us to pick up libpython at runtime.
    lib.linker_allow_shlib_undefined = true;
    lib.addIncludePath(.{ .path = "includes/python" });
    b.installArtifact(lib);
}

main.zig

const std = @import("std");
const Allocator = std.mem.Allocator;

const py = @cImport({
    @cDefine("Py_LIMITED_API", "0x03090000"); // 3.9
    @cInclude("Python.h");
});

export fn hello(self: [*c]py.PyObject, args: [*c]py.PyObject) [*c]py.PyObject {
    _ = args;
    _ = self;
    return py.PyUnicode_FromString("Hello, World!");
}

fn create(allocator: Allocator) !?*py.PyObject {
    const methods: []py.PyMethodDef = try allocator.alloc(py.PyMethodDef, 2);
    methods[0] = py.PyMethodDef{
        .ml_name = "hello",
        .ml_meth = &hello,
        .ml_flags = py.METH_NOARGS,
        .ml_doc = null,
    };
    methods[1] = py.PyMethodDef{
        .ml_name = null,
        .ml_meth = null,
        .ml_flags = 0,
        .ml_doc = null,
    };

    const mod = try allocator.create(py.PyModuleDef);
    mod.* = py.PyModuleDef{
        .m_base = py.PyModuleDef_Base{
            .ob_base = py.PyObject{
                .ob_refcnt = 1,
                .ob_type = null,
            },
            .m_init = null,
            .m_index = 0,
            .m_copy = null,
        },
        .m_name = "mymodule",
        .m_doc = "my docstring",
        .m_size = -1,
        .m_methods = @ptrCast(methods),
        .m_slots = null,
        .m_traverse = null,
        .m_clear = null,
        .m_free = null,
    };

    return py.PyModule_Create(mod);
}

export fn PyInit_libenc() ?*py.PyObject {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    const allocator = gpa.allocator();
    return create(allocator) catch unreachable;
}

pyproject.toml

[build-system]
requires = ["poetry-core", "setuptools"]
build-backend = "poetry.core.masonry.api"

[tool.poetry.build]
script = "build.py"
generate-setup-file = true

[tool.poetry]
...

setuptools_zig.py

import os
import pathlib
import shutil
import subprocess
import sysconfig

from setuptools.command.build_ext import build_ext
from setuptools.extension import Extension


class ZigExtension(Extension):
    def __init__(
        self, name: str, build_zig: str = "build.zig", release_mode: str = "ReleaseFast", py_limited_api=True, **kwargs
    ):
        self.build_zig = pathlib.Path(build_zig)
        self.release_mode = release_mode
        super().__init__(name, sources=[], **kwargs)


class zig_build_ext(build_ext):
    def build_extension(self, ext) -> None:
        if not isinstance(ext, ZigExtension):
            return super().build_extension(ext)

        # Clean the zig-out directory so that we ensure we copy the correct library
        lib_out = ext.build_zig.parent / "zig-out" / "lib"
        shutil.rmtree(lib_out, ignore_errors=True)

        # Setup a symlink to the Python include directory
        includes = ext.build_zig.parent / "includes"
        includes.mkdir(exist_ok=True)
        python_include = includes / "python"
        python_include.unlink(missing_ok=True)
        os.symlink(sysconfig.get_path("include"), str(python_include))

        argv = [
            "zig",
            "build",
            "--build-file",
            ext.build_zig,
            # TODO(ngates): what's the debug flag we should switch on? Inplace maybe?
            f"-Doptimize={'Debug' if self.debug else 'ReleaseFast'}",
        ]

        # Compile our Zig library
        subprocess.run(argv, check=True)

        # Copy it into place
        lib_files = list(lib_out.glob("*"))
        if len(lib_files) != 1:
            raise ValueError(f"Found multiple zig output libraries {lib_files}")
        shutil.copy(lib_files[0], self.get_ext_fullpath(ext.name))