Merge pull request #44 from robotpy/fixed-size-span

Add support for fixed/dynamic sized spans

Author: virtuald

tests/cpp/wpiutil_test/module.cpp

@@ -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;
@@ -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);

tests/test_span.py

@@ -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])

wpiutil/src/type_casters/wpi_span_type_caster.h

@@ -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) {
@@ -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;
   }
 
@@ -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))
@@ -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;
   }
 
@@ -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) {