Skip to content
This repository was archived by the owner on Nov 1, 2023. It is now read-only.

Add support for fixed/dynamic sized spans #44

Merged
merged 1 commit into from
Dec 27, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions tests/cpp/wpiutil_test/module.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,11 @@ std::span<std::vector<std::string>> load_span_vector(std::span<std::vector<std::
return ref;
}

std::span<const double, 3> load_span_fixed_double(std::span<const double, 3> ref) {
return ref;
}


std::span<int> cast_span() {
static std::vector<int> vec{1, 2, 3};
return vec;
Expand Down Expand Up @@ -148,6 +153,7 @@ RPYBUILD_PYBIND11_MODULE(m) {
// span
m.def("load_span_int", &load_span_int);
m.def("load_span_bool", &load_span_bool);
m.def("load_span_fixed_double", &load_span_fixed_double);
m.def("load_span_string", &load_span_string);
m.def("load_span_string_const", &load_span_string_const);
m.def("load_span_string_view", &load_span_string_view);
Expand Down
7 changes: 7 additions & 0 deletions tests/test_span.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,3 +66,10 @@ def test_span_cast():

def test_string_span():
assert module.cast_string_span() == ["hi", "there"]


def test_fixed_double_span():
assert module.load_span_fixed_double([1, 2, 3]) == (1, 2, 3)

with pytest.raises(TypeError):
assert module.load_span_fixed_double([1, 2, 3, 4])
83 changes: 76 additions & 7 deletions wpiutil/src/type_casters/wpi_span_type_caster.h
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,77 @@
namespace pybind11 {
namespace detail {

template <typename Type> struct type_caster<std::span<Type>> {
template <size_t N>
struct span_name_maker {
template <typename T>
static constexpr auto make(const T &t) {
return concat(t, span_name_maker<N-1>::make(t));
}
};

template <>
struct span_name_maker<1> {
template <typename T>
static constexpr auto make(const T &t) {
return t;
}
};

// span with fixed size converts to a tuple
template <typename Type, size_t Extent> struct type_caster<std::span<Type, Extent>> {
using span_type = typename std::span<Type, Extent>;
using value_conv = make_caster<Type>;
using value_type = typename std::remove_cv<Type>::type;

value_type backing_array[Extent] = {};

PYBIND11_TYPE_CASTER(span_type, _("Tuple[") + span_name_maker<Extent>::make(value_conv::name) + _("]"));

type_caster() : value(backing_array) {}

bool load(handle src, bool convert) {
if (!isinstance<sequence>(src) || isinstance<str>(src))
return false;
auto s = reinterpret_borrow<sequence>(src);
if (s.size() != Extent)
return false;
size_t i = 0;
for (auto it : s) {
value_conv conv;
if (!conv.load(it, convert))
return false;
backing_array[i] = cast_op<Type &&>(std::move(conv));
i++;
}
return true;
}

public:
template <typename T>
static handle cast(T &&src, return_value_policy policy, handle parent) {
if (!std::is_lvalue_reference<T>::value)
policy = return_value_policy_override<Type>::policy(policy);
tuple l(Extent);
size_t index = 0;
for (auto &&value : src) {
auto value_ = reinterpret_steal<object>(
value_conv::cast(forward_like<T>(value), policy, parent));
if (!value_)
return handle();
PyTuple_SET_ITEM(l.ptr(), (ssize_t)index++,
value_.release().ptr()); // steals a reference
}
return l.release();
}
};


// span with dynamic extent
template <typename Type> struct type_caster<std::span<Type, std::dynamic_extent>> {
using span_type = typename std::span<Type, std::dynamic_extent>;
using value_conv = make_caster<Type>;
using value_type = typename std::remove_cv<Type>::type;
PYBIND11_TYPE_CASTER(std::span<Type>, _("List[") + value_conv::name + _("]"));
PYBIND11_TYPE_CASTER(span_type, _("List[") + value_conv::name + _("]"));

wpi::SmallVector<value_type, 32> vec;
bool load(handle src, bool convert) {
Expand All @@ -27,7 +94,7 @@ template <typename Type> struct type_caster<std::span<Type>> {
return false;
vec.push_back(cast_op<Type &&>(std::move(conv)));
}
value = std::span<Type>(std::data(vec), std::size(vec));
value = span_type(std::data(vec), std::size(vec));
return true;
}

Expand All @@ -51,8 +118,9 @@ template <typename Type> struct type_caster<std::span<Type>> {
};

// span specialization: accepts any readonly buffers
template <> struct type_caster<std::span<const uint8_t>> {
PYBIND11_TYPE_CASTER(std::span<const uint8_t>, _("buffer"));
template <> struct type_caster<std::span<const uint8_t, std::dynamic_extent>> {
using span_type = typename std::span<const uint8_t, std::dynamic_extent>;
PYBIND11_TYPE_CASTER(span_type, _("buffer"));

bool load(handle src, bool convert) {
if (!isinstance<buffer>(src))
Expand All @@ -63,7 +131,7 @@ template <> struct type_caster<std::span<const uint8_t>> {
return false;
}

value = std::span<const uint8_t>((const uint8_t*)req.ptr, req.size*req.itemsize);
value = span_type((const uint8_t*)req.ptr, req.size*req.itemsize);
return true;
}

Expand All @@ -75,7 +143,8 @@ template <> struct type_caster<std::span<const uint8_t>> {
};

// span specialization: writeable buffer
template <> struct type_caster<std::span<uint8_t>> {
template <> struct type_caster<std::span<uint8_t, std::dynamic_extent>> {
using span_type = typename std::span<const uint8_t, std::dynamic_extent>;
PYBIND11_TYPE_CASTER(std::span<uint8_t>, _("buffer"));

bool load(handle src, bool convert) {
Expand Down