Skip to content

Python Buffer Protocol

Python objects implementing the Buffer Protocol can be used with zero copy.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
pub fn sum(args: struct { buf: py.PyObject }) !i64 {
    const view = try args.buf.getBuffer(py.PyBuffer.Flags.ND);
    defer view.release();

    var bufferSum: i64 = 0;
    for (view.asSlice(i64)) |value| bufferSum += value;
    return bufferSum;
}

comptime {
    py.rootmodule(@This());
}

This function accepts Python arrays, Numpy arrays, or any other buffer protocol implementation.

1
2
3
4
5
def test_sum():
    import numpy as np

    arr = np.array([1, 2, 3, 4, 5], dtype=np.int64)
    assert buffers.sum(arr) == 15

Note

Understanding request types is important when working with buffers. Common request types are implemented as py.PyBuffer.Flags, e.g. py.PyBuffer.Flags.FULL_RO.

You can implement a buffer protocol in a Pydust module by implementing __buffer__ and optionally __release_buffer__ methods.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
pub const ConstantBuffer = py.class(struct {
    pub const __doc__ = "A class implementing a buffer protocol";
    const Self = @This();

    values: []i64,
    shape: []const isize, // isize to be compatible with Python API
    format: [:0]const u8 = "l", // i64

    pub fn __init__(self: *Self, args: struct { elem: i64, length: u32 }) !void {
        const values = try py.allocator.alloc(i64, args.length);
        @memset(values, args.elem);

        const shape = try py.allocator.alloc(isize, 1);
        shape[0] = @intCast(args.length);

        self.* = .{ .values = values, .shape = shape };
    }

    pub fn __del__(self: *Self) void {
        py.allocator.free(self.values);
        py.allocator.free(self.shape);
    }

    pub fn __buffer__(self: *const Self, view: *py.PyBuffer, flags: c_int) !void {
        // For more details on request types, see https://docs.python.org/3/c-api/buffer.html#buffer-request-types
        if (flags & py.PyBuffer.Flags.WRITABLE != 0) {
            return py.BufferError.raise("request for writable buffer is rejected");
        }
        view.initFromSlice(i64, self.values, self.shape, self);
    }
});