Search code examples
pythonc++swigswig-typemap

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


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())

Solution

  • 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)