Support Buffer Protocol

example/buffers.zig

@@ -0,0 +1,67 @@
+const std = @import("std");
+const py = @import("pydust");
+
+pub const ConstantBuffer = py.class("ConstantBuffer", struct {
+    pub const __doc__ = "A class implementing a buffer protocol";
+    const Self = @This();
+
+    values: []i64,
+
+    pub fn __init__(self: *Self, args: *const extern struct { elem: py.PyLong, size: py.PyLong }) !void {
+        const elem = try args.elem.as(i64);
+        const size = try args.size.as(u64);
+
+        self.values = try py.allocator.alloc(i64, size);
+        @memset(self.values, elem);
+    }
+
+    // TODO(marko): Get obj from self.
+    pub fn __buffer__(self: *const Self, obj: py.PyObject, view: *py.PyBuffer, flags: c_int) !void {
+        if (flags & py.ffi.PyBUF_WRITABLE != 0) {
+            return py.BufferError.raise("Must not request writable");
+        }
+
+        const shape = try py.allocator.alloc(isize, 1);
+        shape[0] = @intCast(self.values.len);
+
+        // Because we're using values, we need to incref it.
+        obj.incref();
+
+        view.* = .{
+            .buf = std.mem.sliceAsBytes(self.values).ptr,
+            .obj = obj.py,
+            .len = @intCast(self.values.len * @sizeOf(i64)),
+            .readonly = 1,
+            .itemsize = @sizeOf(i64),
+            .format_str = try py.PyBuffer.allocFormat(i64, py.allocator),
+            .ndim = 1,
+            .shape = shape.ptr,
+            .strides = null,
+            .suboffsets = null,
+            .internal = null,
+        };
+    }
+
+    pub fn __release_buffer__(self: *const Self, view: *py.PyBuffer) void {
+        py.allocator.free(self.values);
+        py.allocator.free(view.format_str[0..@intCast(std.mem.indexOfSentinel(u8, 0, view.format_str) + 1)]);
+        if (view.shape) |shape| py.allocator.free(shape[0..@intCast(view.ndim)]);
+        view.obj = null;
+    }
+});
+
+// Accept a buffer protocol object.
+pub fn sum(args: *const extern struct { buf: py.PyObject }) !py.PyLong {
+    var view = try py.PyBuffer.of(args.buf, py.ffi.PyBUF_C_CONTIGUOUS);
+    defer view.release();
+
+    const sliceView: []i64 = @alignCast(std.mem.bytesAsSlice(i64, view.buf.?[0..@intCast(view.len)]));
+    var bufferSum: i64 = 0;
+    for (sliceView) |value| bufferSum += value;
+
+    return try py.PyLong.from(i64, bufferSum);
+}
+
+comptime {
+    py.module(@This());
+}

pydust/src/errors.zig

@@ -1,8 +1,8 @@
-const std = @import("std");
+const Allocator = @import("std").mem.Allocator;
 
 pub const PyError = error{
     // Propagate an error raised from another Python function call.
     // This is the equivalent of returning PyNULL and allowing the already set error info to remain.
     Propagate,
     Raised,
-} || std.mem.Allocator.Error;
+} || Allocator.Error;

pydust/src/functions.zig

@@ -19,6 +19,8 @@ const reservedNames = .{
     "__new__",
     "__init__",
     "__del__",
+    "__buffer__",
+    "__release_buffer__",
 };
 
 /// Parse the arguments of a Zig function into a Pydust function siganture.

pydust/src/pytypes.zig

@@ -91,6 +91,20 @@ fn Slots(comptime name: [:0]const u8, comptime definition: type, comptime Instan
                 }};
             }
 
+            if (@hasDecl(definition, "__buffer__")) {
+                slots_ = slots_ ++ .{ffi.PyType_Slot{
+                    .slot = ffi.Py_bf_getbuffer,
+                    .pfunc = @ptrCast(@constCast(&bf_getbuffer)),
+                }};
+            }
+
+            if (@hasDecl(definition, "__release_buffer__")) {
+                slots_ = slots_ ++ .{ffi.PyType_Slot{
+                    .slot = ffi.Py_bf_releasebuffer,
+                    .pfunc = @ptrCast(@constCast(&bf_releasebuffer)),
+                }};
+            }
+
             slots_ = slots_ ++ .{ffi.PyType_Slot{
                 .slot = ffi.Py_tp_methods,
                 .pfunc = @ptrCast(@constCast(&methods.pydefs)),
@@ -119,6 +133,19 @@ fn Slots(comptime name: [:0]const u8, comptime definition: type, comptime Instan
 
             ffi.PyErr_Restore(error_type, error_value, error_tb);
         }
+
+        fn bf_getbuffer(self: *ffi.PyObject, view: *ffi.Py_buffer, flags: c_int) callconv(.C) c_int {
+            // In case of any error, the view.obj field must be set to NULL.
+            view.obj = null;
+
+            const instance: *Instance = @ptrCast(self);
+            return tramp.errVoid(definition.__buffer__(&instance.state, .{ .py = self }, @ptrCast(view), flags));
+        }
+
+        fn bf_releasebuffer(self: *ffi.PyObject, view: *ffi.Py_buffer) callconv(.C) void {
+            const instance: *Instance = @ptrCast(self);
+            return definition.__release_buffer__(&instance.state, @ptrCast(view));
+        }
     };
 }
 

pydust/src/types.zig

@@ -1,4 +1,5 @@
 pub usingnamespace @import("types/bool.zig");
+pub usingnamespace @import("types/buffer.zig");
 pub usingnamespace @import("types/dict.zig");
 pub usingnamespace @import("types/error.zig");
 pub usingnamespace @import("types/float.zig");

pydust/src/types/buffer.zig

@@ -0,0 +1,109 @@
+const std = @import("std");
+const py = @import("../pydust.zig");
+const ffi = py.ffi;
+const PyError = @import("../errors.zig").PyError;
+
+/// Wrapper for Python Py_buffer.
+/// See: https://docs.python.org/3/c-api/buffer.html
+pub const PyBuffer = extern struct {
+    const Self = @This();
+
+    buf: ?[*]u8,
+
+    // TODO(marko): We can make this PyObject but have to make ffi reference optional.
+    obj: ?*ffi.PyObject,
+
+    // product(shape) * itemsize.
+    // For contiguous arrays, this is the length of the underlying memory block.
+    // For non-contiguous arrays, it is the length that the logical structure would
+    // have if it were copied to a contiguous representation.
+    len: isize,
+    itemsize: isize,
+    readonly: c_int,
+
+    // If ndim == 0, the memory location pointed to by buf is interpreted as a scalar of size itemsize.
+    // In that case, both shape and strides are NULL.
+    ndim: c_int,
+    format_str: [*:0]u8,
+
+    shape: ?[*]isize,
+    // If strides is NULL, the array is interpreted as a standard n-dimensional C-array.
+    // Otherwise, the consumer must access an n-dimensional array as follows:
+    // ptr = (char *)buf + indices[0] * strides[0] + ... + indices[n-1] * strides[n-1];
+    strides: ?[*]isize,
+    // If all suboffsets are negative (i.e. no de-referencing is needed),
+    // then this field must be NULL (the default value).
+    suboffsets: ?[*]isize,
+
+    internal: ?*anyopaque,
+
+    pub fn release(self: *Self) void {
+        ffi.PyBuffer_Release(@ptrCast(self));
+    }
+
+    pub fn pyObj(self: *Self) py.PyObject {
+        return .{ .py = self.obj orelse unreachable };
+    }
+
+    // Flag is a combination of ffi.PyBUF_* flags.
+    // See: https://docs.python.org/3/c-api/buffer.html#buffer-request-types
+    pub fn of(obj: py.PyObject, flag: c_int) !PyBuffer {
+        if (ffi.PyObject_CheckBuffer(obj.py) != 1) {
+            return py.BufferError.raise("object does not support buffer interface");
+        }
+
+        var out: Self = undefined;
+        if (ffi.PyObject_GetBuffer(obj.py, @ptrCast(&out), flag) != 0) {
+            // Error is already raised.
+            return PyError.Propagate;
+        }
+        return out;
+    }
+
+    // A helper function for converting Zig types to buffer format string.
+    pub fn allocFormat(comptime value_type: type, allocator: std.mem.Allocator) ![*:0]u8 {
+        const fmt = PyBuffer.getFormat(value_type);
+        var fmt_c = try allocator.allocSentinel(u8, fmt.len, 0);
+        @memcpy(fmt_c, fmt);
+        return fmt_c;
+    }
+
+    fn getFormat(comptime value_type: type) []const u8 {
+        switch (@typeInfo(value_type)) {
+            .Int => |i| {
+                switch (i.signedness) {
+                    .unsigned => switch (i.bits) {
+                        8 => return &.{'B'},
+                        16 => return &.{'H'},
+                        32 => return &.{'I'},
+                        64 => return &.{'L'},
+                        else => {
+                            @compileError("Unsupported buffer value type" ++ @typeName(value_type));
+                        },
+                    },
+                    .signed => switch (i.bits) {
+                        8 => return &.{'b'},
+                        16 => return &.{'h'},
+                        32 => return &.{'i'},
+                        64 => return &.{'l'},
+                        else => {
+                            @compileError("Unsupported buffer value type" ++ @typeName(value_type));
+                        },
+                    },
+                }
+            },
+            .Float => |f| {
+                switch (f.bits) {
+                    32 => return &.{'f'},
+                    64 => return &.{'d'},
+                    else => {
+                        @compileError("Unsupported buffer value type" ++ @typeName(value_type));
+                    },
+                }
+            },
+            else => {
+                @compileError("Unsupported buffer value type" ++ @typeName(value_type));
+            },
+        }
+    }
+};

pyproject.toml

@@ -81,3 +81,7 @@ root = "example/functions.zig"
 [[tool.pydust.ext_module]]
 name = "example.classes"
 root = "example/classes.zig"
+
+[[tool.pydust.ext_module]]
+name = "example.buffers"
+root = "example/buffers.zig"

test/test_buffers.py

@@ -0,0 +1,17 @@
+from array import array
+
+from example import buffers
+
+
+def test_view():
+    buffer = buffers.ConstantBuffer(1, 10)
+    view = memoryview(buffer)
+    for i in range(10):
+        assert view[i] == 1
+    view.release()
+
+
+def test_sum():
+    # array implements a buffer protocol
+    arr = array("l", [1, 2, 3, 4, 5])
+    assert buffers.sum(arr) == 15