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