SWIG: pass an std::span as argument to python function

176 views Asked by At

I need some help with SWIG on this.

I have a C++ pure virtual function that I need to implement in python and it gets an std:span<uint8_t> as an output argument.

virtual void fill_buffer(size_t offset, std::span<uint8_t> buffer) = 0;

The python side is responsible to fill up that buffer with data. What I did so far to make it work is create some utility functions like:

%inline %{
    PyObject* size(const std::span<uint8_t>& span) {
        return PyInt_FromLong(span.size());
    }

    void fill_buffer(const std::span<uint8_t>& span, const std::vector<uint8_t>& buffer) {
        std::copy(buffer.begin(), buffer.end(), span.data());
    }
%}

And then in the python side I have:

def fill_buffer(self, offset, buffer):
    buffer_size = size(buffer)
    with open(self.resource_file, 'rb') as file:
        file.seek(offset)
        read_bytes = file.read(buffer_size)
    fill_buffer(buffer, read_bytes)

But I am thinking there must be a better way to do this. Maybe using a typemap? I would like to seamlessly use the buffer object in python without the helper functions, maybe something like:

def fill_buffer(self, offset, buffer):
    with open(self.resource_file, 'rb') as file:
        file.seek(offset)
        buffer = file.read(buffer.size())
1

There are 1 answers

0
Flexo On BEST ANSWER

Given the following, complete test:

#include <span>
#include <iostream>

struct test_base {
    virtual void fill_buffer(size_t offset, std::span<uint8_t> buffer) = 0;
    virtual ~test_base() {}
};

inline void run_test(test_base& tb) {
    uint8_t buffer[1024];
    // call virtual function and just print output to prove it worked
    tb.fill_buffer(0, std::span{buffer});
    std::cout << "Buffer is: " << buffer << "\n";
}

Our goal is to wrap it nicely into Python, such that the following example can work:

import test

class PythonSpan(test.test_base):
    def __init__(self):
        super().__init__()

    def fill_buffer(self, offset, buf):
        # totally ignored offset param
        with open(__file__, 'rb') as f:
            got = f.readinto(buf)
            buf[got] = 0 # null terminate for demo

filler = PythonSpan()
test.run_test(filler)

(Note that readinto is the neat way to go from file read straight into a buffer of your choosing)

What we need to make this work is a "directorin" typemap. As luck would have it Python 3's C API has a function that does almost exactly what we want. PyMemoryView_FromMemory creates a buffer object that's a pretty good Python equivalent for std::span.

%module(directors="1") test
%feature("director");

%typemap(directorin) std::span %{
        $input = PyMemoryView_FromMemory(reinterpret_cast<char*>($1.data()), $1.size(), PyBUF_WRITE);
%}

%{
#include "test.h"
%}

%include "test.h"

We can do fancier things if the type in your span was not just a uint8_t, but for this case the simple interface is sufficient. (Also since the buffer we create allows for in place modification there's no need for a "directorargout" typemap, however an "in" typemap might be a useful addition here)