Bindings for std::span

[Walter-Reactor]

Has there been any consideration to adding bindings for std::span or other span-like objects? I'm trying to write a binding and finding it's a challenge to juggle the lifetimes, and wondering if there's any prior art here. list_binding sadly doesn't seem to work well.

Going from python, it seems like there wouldn't be any real way to indicate that C++ owns the data without some kind of wrapper around the span object to free the refcount when the object is destroyed, and going from C++ seems even more dangerous unless you go through and immediately create castings/clones of every held object. Am I missing anything here?

[wjakob]

Nanobind is feature-scoped to C++17. It might make sense to develop some kind of C++20 (and beyond) extension layer, but I do not wish to have such functionality as part of this project.

[Walter-Reactor]

I had the same questions. I wound up with this which has so far worked with the very simple test cases I've thrown at it - It's a modified version of some bindings for abseil::Span I found

template <typename T>
    class type_caster<span<T>>
    {
        // The type referenced by the span, with const removed.
        using value_type = std::remove_cv_t<T>;
        // Avoid pitfalls with std::vector<bool>
        using ListCaster = list_caster<std::vector<std::conditional_t<std::is_same_v<bool, value_type>, uint8_t, value_type>>, value_type>;
        static_assert(sizeof(bool) == sizeof(uint8_t), "bool representation size is unexpected!");

        std::optional<ListCaster> list_caster;

    public:
        using Value = span<T>;
        Value value = span<T>(nullptr, 0);

        static constexpr auto Name = ListCaster::Name;

        // We do not allow moving because 1) spans are super lightweight, so there's
        // no advantage to moving and 2) the span cannot exist without the caster,
        // so moving leaves an implicit dependency (while a reference or pointer
        // make that dependency explicit).
        template <typename T_>
        using Cast = movable_cast_t<T_>;
        operator Value *() { return &value; }
        operator Value &() { return value; }

        // Cast Python -> C++ (nb::cast call)
        bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept
        {
            uint8_t local_flags = flags_for_local_caster<T>(flags);

            // First try to convert from ndarray for efficiency
            // Try to get a 1D contiguous array directly using type tags
            if constexpr (is_ndarray_scalar_v<T>)
            {
                auto caster = make_caster<nb::ndarray<T, ndim<1>>>();
                if (caster.from_python(src, local_flags, cleanup))
                {
                    // Create a span from the array data
                    value = span<T>(caster.value.data(), caster.value.shape(0));
                    return true;
                }
            }

            // Attempt to convert a native sequence. If the is_base_of_v check passes,
            // the elements do not require converting and pointers do not reference a
            // temporary object owned by the element caster. Pointers to converted
            // types are not allowed because they would result a dangling reference
            // when the element caster is destroyed.
            if (std::is_const_v<T> &&
                (!std::is_pointer_v<T> || is_base_caster_v<make_caster<T>>))
            {
                list_caster.emplace();
                if (list_caster->from_python(src, local_flags, cleanup))
                {
                    // Reinterpret cast here is required to handle vector<bool>
                    value = span<T>(reinterpret_cast<T *>(list_caster->value.data()), list_caster->value.size());
                    return true;
                }
                else
                {
                    list_caster.reset();
                }
            }

            return false; // Python type cannot be loaded into a span.
        }

        // Cast C++ -> Python (when returning a span from a C++ function)
        static handle from_cpp(span<T> src, rv_policy policy, cleanup_list *cleanup)
        {
            return ListCaster::from_cpp(src, policy, cleanup);
        }
    };

[bethebunny]

Thanks for the start :) For from_cpp, it seems like rv_policy::reference_internal should be useful in safely covering cases where you want to return a view to data in a method. I'm trying to figure out if I can override the policy in the caster, or set a different default policy for spans.

Edit: actually basically what I want is close to make_iterator, but I'd want it to always nb::keep_alive on the parent object.

[Walter-Reactor]

Unfortunately the approach I was taking runs into issues if you have a std::optionalstd::span. The Caster may be constructed as part of another caster if you have a std::optional span or worse, a std::vector of span, so having the list_caster live for the duration of the main caster leads to dangling memory. Instead, I simply move the std::vector onto the heap, stick it in a capsule, and put it in the cleanup list. This seems to work in the testing I've done thus far, I'll post if I find any other problems.

[bethebunny]

Here's the janky thing I've landed on for now: https://gist.github.com/bethebunny/0d38b8cb45a3586c567b7f37d969c2b0 lots to improve but it is getting the job done for the time being :)

[Walter-Reactor]

ahh, that solution with specializing the outer type_caster is quite clever! I'm sticking with making the capsule store things for now since it's more general, though does lead to a bit of overhead.

[lyd405121]

Share My Code

• I implement a scene graph library to load openusd, gltf, mjcf,also mitsuba xml
• And using std::span to import/export textures, geometry, instance,camera,lights, materials in blender as a plugin
• Here is my header,very ugly, but straightforward

#pragma once
#include "nb_span.h"
NAMESPACE_BEGIN(NB_NAMESPACE)
NAMESPACE_BEGIN(detail)


template <>
struct type_caster<std::span<unsigned char>>
    : span_numpy_caster<unsigned char> {
};

template <>
struct type_caster<std::span<char>>
    : span_numpy_caster<char> {
};

template <>
struct type_caster<std::span<int>>
    : span_numpy_caster<int> {
};

template <>
struct type_caster<std::span<unsigned int>>
    : span_numpy_caster<unsigned int> {
};

template <>
struct type_caster<std::span<size_t>>
    : span_numpy_caster<size_t> {
};

template <>
struct type_caster<std::span<float>>
    : span_numpy_caster<float> {
};

template <>
struct type_caster<std::span<double>>
    : span_numpy_caster<double> {
};


template <>
struct type_caster<std::span<HKMATH::HKVector2>>
    : span_numpy_vec_caster<HKMATH::HKVector2, float, 2> {
};

template <>
struct type_caster<std::span<HKMATH::HKVector3>>
    : span_numpy_vec_caster<HKMATH::HKVector3, float, 3> {
};

template <>
struct type_caster<std::span<HKMATH::HKVector4>>
    : span_numpy_vec_caster< HKMATH::HKVector4, float, 4> {
};

template <>
struct type_caster<std::span<HKMATH::HKVector4I>>
    : span_numpy_vec_caster<HKMATH::HKVector4I, int, 4> {
};

template <>
struct type_caster<std::span<HKMATH::HKMatrix>>
    : span_numpy_mat_caster< HKMATH::HKMatrix, float, 4, 4> {
};

template <typename ElementType> struct type_caster<std::span<ElementType>>
 : span_list_caster<std::span<ElementType>, ElementType> { };

NAMESPACE_END(detail)
NAMESPACE_END(NB_NAMESPACE)

• I divide different data structure into three kinds of macro, basic data, math data, class data
• Then cames the impl:

#pragma once

#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
#include <span>
NAMESPACE_BEGIN(NB_NAMESPACE)
NAMESPACE_BEGIN(detail)

template <typename T>
struct span_numpy_caster {
    using Span = std::span<T>;
    using value_type = T;
    using NDArray = ndarray<T, numpy>;
    using NDArrayCaster = make_caster<NDArray>;

    NB_TYPE_CASTER(Span, NDArrayCaster::Name)

        bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept {
        make_caster<NDArray> caster;
        if (!caster.from_python(src, flags, cleanup))
            return false;
        const NDArray& array = caster.value;
        value = Span((T*)array.data(), array.size());
        return true;
    }

    static handle from_cpp(const Span& v, rv_policy policy, cleanup_list* cleanup) noexcept {
        if (policy == rv_policy::automatic || policy == rv_policy::automatic_reference)
            policy = rv_policy::reference;
        size_t shape = v.size();
        void* data_ptr = const_cast<T*>(v.data());
        object owner;
        switch (policy) {
        case rv_policy::move:
            owner = capsule(new Span(std::move(v)), [](void* p) noexcept { delete (Span*)p; });
            data_ptr = ((Span*)owner.ptr())->data();
            policy = rv_policy::reference;
            break;
        case rv_policy::reference_internal:
            if (cleanup->self()) owner = borrow(cleanup->self());
            policy = rv_policy::reference;
            break;
        default: break;
        }
        NDArray array(data_ptr, 1, &shape, owner, nullptr);
        return NDArrayCaster::from_cpp(array, policy, cleanup);
    }
};

// 支持 std::span<Vector>，如 Vector3(float[3])，底层数据是 float
template <typename VecType, typename ScalarType, size_t Dim>
struct span_numpy_vec_caster {
    using Span = std::span<VecType>;
    using NDArray = ndarray<ScalarType, numpy>;
    using NDArrayCaster = make_caster<NDArray>;

    NB_TYPE_CASTER(Span, NDArrayCaster::Name)

        bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept {
        make_caster<NDArray> caster;
        if (!caster.from_python(src, flags, cleanup))
            return false;
        const NDArray& array = caster.value;
        if (array.ndim() != 2 || array.shape(1) != Dim)
            return false;
        const ScalarType* data = static_cast<const ScalarType*>(array.data());
        size_t n = array.shape(0);
        value = Span((VecType*)(data), n);
        return true;
    }

    static handle from_cpp(const Span& v, rv_policy policy, cleanup_list* cleanup) noexcept {
        if (policy == rv_policy::automatic || policy == rv_policy::automatic_reference)
            policy = rv_policy::reference;
        size_t shape[2] = { v.size(), Dim };
        ScalarType* data = const_cast<ScalarType*>(reinterpret_cast<const ScalarType*>(v.data()));
        object owner;
        switch (policy) {
        case rv_policy::move:
            owner = capsule(new Span(v), [](void* p) noexcept { delete (Span*)p; });
            data = reinterpret_cast<ScalarType*>(((Span*)owner.ptr())->data());
            policy = rv_policy::reference;
            break;
        case rv_policy::reference_internal:
            if (cleanup->self()) owner = borrow(cleanup->self());
            policy = rv_policy::reference;
            break;
        default: break;
        }
        NDArray array(data, 2, shape, owner, nullptr);
        return NDArrayCaster::from_cpp(array, policy, cleanup);
    }
};

// 支持 std::span<Mat>，如 Mat4x4(float[16])，底层数据是 float
template <typename MatType, typename ScalarType, size_t Rows, size_t Cols>
struct span_numpy_mat_caster {
    using Span = std::span<MatType>;
    using NDArray = ndarray<ScalarType, numpy>;
    using NDArrayCaster = make_caster<NDArray>;

    NB_TYPE_CASTER(Span, NDArrayCaster::Name)

        bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept {
        make_caster<NDArray> caster;
        if (!caster.from_python(src, flags, cleanup))
            return false;
        const NDArray& array = caster.value;
        if (array.ndim() != 3 || array.shape(1) != Rows || array.shape(2) != Cols)
            return false;
        const ScalarType* data = static_cast<const ScalarType*>(array.data());
        size_t n = array.shape(0);
        value = Span((MatType*)(data), n);
        return true;
    }

    static handle from_cpp(const Span& v, rv_policy policy, cleanup_list* cleanup) noexcept {
        if (policy == rv_policy::automatic || policy == rv_policy::automatic_reference)
            policy = rv_policy::reference;
        size_t shape[3] = { v.size(), Rows, Cols };
        ScalarType* data = const_cast<ScalarType*>(reinterpret_cast<const ScalarType*>(v.data()));
        object owner;
        switch (policy) {
        case rv_policy::move:
            owner = capsule(new Span(v), [](void* p) noexcept { delete (Span*)p; });
            data = reinterpret_cast<ScalarType*>(((Span*)owner.ptr())->data());
            policy = rv_policy::reference;
            break;
        case rv_policy::reference_internal:
            if (cleanup->self()) owner = borrow(cleanup->self());
            policy = rv_policy::reference;
            break;
        default: break;
        }
        NDArray array(data, 3, shape, owner, nullptr);
        return NDArrayCaster::from_cpp(array, policy, cleanup);
    }
};

std::unordered_map<void*, void*> g_mapPySpanData;
template <typename Entry>
Entry* NewSpanDataFor(void* pPyObject, size_t size)
{
    auto it = g_mapPySpanData.find(pPyObject);
    if (it != g_mapPySpanData.end())
    {
        delete[] it->second;
        g_mapPySpanData.erase(it);
    }
    Entry* pSpanData = new Entry[size];
    g_mapPySpanData.insert({ pPyObject, pSpanData });
    return pSpanData;
}

template <typename List, typename Entry> struct span_list_caster {
    NB_TYPE_CASTER(List, io_name(NB_TYPING_SEQUENCE, NB_TYPING_LIST) +
        const_name("[") + make_caster<Entry>::Name +
        const_name("]"))
    
    using Caster = make_caster<Entry>;

    template <typename T> using has_reserve = decltype(std::declval<T>().reserve(0));

    bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept {
        size_t size;
        PyObject* temp;

        /* Will initialize 'size' and 'temp'. All return values and
           return parameters are zero/NULL in the case of a failure. */
        PyObject** o = seq_get(src.ptr(), &size, &temp);

        Entry* pNewData = NewSpanDataFor<Entry>(o, size);

        Caster caster;
        bool success = o != nullptr;

        flags = flags_for_local_caster<Entry>(flags);

        for (size_t i = 0; i < size; ++i) {
            if (!caster.from_python(o[i], flags, cleanup) ||
                !caster.template can_cast<Entry>()) {
                success = false;
                break;
            }
            pNewData[i] = caster.operator cast_t<Entry>();
        }
        value = std::span<Entry>(pNewData, size);
        Py_XDECREF(temp);

        return success;
    }

    template <typename T>
    static handle from_cpp(T&& src, rv_policy policy, cleanup_list* cleanup) {
        object ret = steal(PyList_New(src.size()));

        if (ret.is_valid()) {
            Py_ssize_t index = 0;

            for (auto&& value : src) {
                handle h = Caster::from_cpp(forward_like_<T>(value), policy, cleanup);

                if (!h.is_valid()) {
                    ret.reset();
                    break;
                }

                NB_LIST_SET_ITEM(ret.ptr(), index++, h.ptr());
            }
        }

        return ret.release();
    }
};

NAMESPACE_END(detail)
NAMESPACE_END(NB_NAMESPACE)

• For basic data and math data, transfer them into numpy array
• For complecated class data, using python list, and make a global memory container to hold(maybe memory leak here)
• I am very basic cpp user, don't laugh at my ugly code, but it works very fast

[lyd405121]

My binding code

• this is the code auto-gen using libclang

    nb::class_<SG_GEOMETRY_VERTEX, SG_OBJECT>(m, "SG_GEOMETRY_VERTEX")
    .def(nb::init<>())
    .def_rw("stAABB", &SG_GEOMETRY_VERTEX::stAABB)
    .def_rw("spanPosition", &SG_GEOMETRY_VERTEX::spanPosition)
    .def_rw("spanNormal", &SG_GEOMETRY_VERTEX::spanNormal)
    .def_rw("spanTexcoord", &SG_GEOMETRY_VERTEX::spanTexcoord)
    .def_rw("spanTangent", &SG_GEOMETRY_VERTEX::spanTangent)
    .def_rw("spanBitangent", &SG_GEOMETRY_VERTEX::spanBitangent)
    .def_rw("spanColor", &SG_GEOMETRY_VERTEX::spanColor)
    ;
		
    nb::class_<SG_GEOMETRY, SG_MEDIA>(m, "SG_GEOMETRY")
    .def(nb::init<>())
    .def_rw("emType", &SG_GEOMETRY::emType)
    .def_rw("emBuiltinType", &SG_GEOMETRY::emBuiltinType)
    .def_rw("bBuildAdjacencyBuffer", &SG_GEOMETRY::bBuildAdjacencyBuffer)
    .def_rw("nIndiceCnt", &SG_GEOMETRY::nIndiceCnt)
    .def_rw("nVerticeCnt", &SG_GEOMETRY::nVerticeCnt)
    .def_rw("spanIndex", &SG_GEOMETRY::spanIndex)
    .def_rw("spanBoneWeight", &SG_GEOMETRY::spanBoneWeight)
    .def_rw("spanBoneIndex", &SG_GEOMETRY::spanBoneIndex)
    .def_rw("stVertex", &SG_GEOMETRY::stVertex)
    .def_rw("pAsset", &SG_GEOMETRY::pAsset)
    ;

• And a skin-animation function in python, lots of math operation with numpy,but very fast

        for ani in delta_frame.spanSkin:
            matrices = ani.spanMatrix  
            if len(ani.spanMatrix) == 0:
                continue
            
            for geo in ani.spanGeometry:
                if len(geo.stVertex.spanPosition) == 0:
                    continue
                positions = geo.stVertex.spanPosition# (N, 3)
                bone_weights = geo.spanBoneWeight     # (N, 4)
                bone_indices = geo.spanBoneIndex       # (N, 4)
                N = positions.shape[0]
                skinned = np.zeros((N, 3), dtype=np.float32)

                # 扩展 positions 为 (N, 4)
                positions_h = np.concatenate([positions, np.ones((N, 1), dtype=positions.dtype)], axis=1)  # (N, 4)

                for j in range(4):
                    vWeight = bone_weights[:, j]      # (N,)
                    vIndex = bone_indices[:, j]       # (N,)
                    valid = (vWeight >= 0.001) & (vIndex >= 0)
                    if not np.any(valid):
                        continue
                    idx = vIndex[valid].astype(np.int32)
                    m = matrices[idx]                 # (n_valid, 4, 4)
                    pos = positions_h[valid][:, None, :]  # (n_valid, 1, 4)
                    # 批量矩阵乘法
                    transformed = np.matmul(m, pos.transpose(0,2,1)).squeeze(-1)  # (n_valid, 4)
                    transformed = transformed[:, :3] * vWeight[valid, None]       # (n_valid, 3)
                    skinned[valid] += transformed

                #print(type(skinned[0]),skinned[0].shape,type(skinned),skinned[0].dtype)
                #print(type(geo.stVertex.spanPosition[0]),geo.stVertex.spanPosition[0].shape,geo.stVertex.spanPosition[0].dtype)
                #print(type(geo.stVertex.spanPosition),len(geo.stVertex.spanPosition))
                geo.stVertex.spanPosition = skinned