Merge pull request #143 from robotpy/struct-array

wpistruct: Add packArray/unpackArray

Author: virtuald

subprojects/robotpy-wpiutil/gen/WPyStruct.yml

@@ -14,12 +14,18 @@ functions:
   pack:
     subpackage: wpistruct
     no_release_gil: true
+  packArray:
+    subpackage: wpistruct
+    no_release_gil: true
   packInto:
     subpackage: wpistruct
     no_release_gil: true
   unpack:
     subpackage: wpistruct
     no_release_gil: true
+  unpackArray:
+    subpackage: wpistruct
+    no_release_gil: true
   unpackInto:
     subpackage: wpistruct
     no_release_gil: true

subprojects/robotpy-wpiutil/tests/test_struct.py

@@ -46,6 +46,10 @@ def test_pack():
     assert wpistruct.pack(module.ThingA(1)) == b"\x01"
 
 
+def test_pack_array():
+    assert wpistruct.packArray([module.ThingA(1), module.ThingA(2)]) == b"\x01\x02"
+
+
 def test_pack_into():
     buf = bytearray(1)
     wpistruct.packInto(module.ThingA(1), buf)
@@ -62,6 +66,13 @@ def test_unpack():
     assert wpistruct.unpack(module.ThingA, b"\x01") == module.ThingA(1)
 
 
+def test_unpack_array():
+    assert wpistruct.unpackArray(module.ThingA, b"\x01\x02") == [
+        module.ThingA(1),
+        module.ThingA(2),
+    ]
+
+
 # def test_unpack_into():
 #     r1 = module.ThingA(1)
 #     r2 = module.ThingA(2)

subprojects/robotpy-wpiutil/wpiutil/src/wpistruct/wpystruct.h

@@ -9,6 +9,7 @@
 #include <wpi/struct/Struct.h>
 
 #include <pybind11/functional.h>
+#include <pybind11/typing.h>
 #include <robotpy_build.h>
 
 static inline std::string pytypename(const py::type &t) {

subprojects/robotpy-wpiutil/wpiutil/src/wpistruct/wpystruct_fns.cpp

@@ -45,6 +45,40 @@ py::bytes pack(const WPyStruct &v) {
   return py::reinterpret_steal<py::bytes>(b);
 }
 
+py::bytes packArray(const py::sequence &seq) {
+  auto len = seq.size();
+  if (len == 0) {
+    return {};
+  }
+
+  WPyStructInfo info(py::type::of(seq[0]));
+  auto sz = wpi::GetStructSize<WPyStruct>(info);
+  auto total = sz*len;
+
+  PyObject *b = PyBytes_FromStringAndSize(NULL, total);
+  if (b == NULL) {
+    throw py::error_already_set();
+  }
+
+  char *pybuf;
+  py::ssize_t pysz;
+  if (PyBytes_AsStringAndSize(b, &pybuf, &pysz) != 0) {
+    Py_DECREF(b);
+    throw py::error_already_set();
+  }
+
+  auto bytes_obj = py::reinterpret_steal<py::bytes>(b);
+
+  for (const auto &v: seq) {
+    WPyStruct wv(v);
+    auto s = std::span((uint8_t *)pybuf, sz);
+    wpi::PackStruct(s, wv, info);
+    pybuf += sz;
+  }
+
+  return bytes_obj;
+}
+
 void packInto(const WPyStruct &v, py::buffer &b) {
   WPyStructInfo info(v);
   py::ssize_t sz = wpi::GetStructSize<WPyStruct>(info);
@@ -83,6 +117,35 @@ WPyStruct unpack(const py::type &t, const py::buffer &b) {
   return wpi::UnpackStruct<WPyStruct, WPyStructInfo>(s, info);
 }
 
+py::typing::List<WPyStruct> unpackArray(const py::type &t, const py::buffer &b) {
+  WPyStructInfo info(t);
+  py::ssize_t sz = wpi::GetStructSize<WPyStruct>(info);
+
+  auto req = b.request();
+  if (req.itemsize != 1) {
+    throw py::value_error("buffer must only contain bytes");
+  } else if (req.ndim != 1) {
+    throw py::value_error("buffer must only have a single dimension");
+  }
+
+  if (req.size % sz != 0) {
+    throw py::value_error("buffer must be multiple of " + std::to_string(sz) + " bytes");
+  }
+
+  auto items = req.size / sz;
+  py::list a(items);
+  const uint8_t *ptr = (const uint8_t *)req.ptr;
+  for (py::ssize_t i = 0; i < items; i++) {
+    auto s = std::span(ptr, sz);
+    auto v = wpi::UnpackStruct<WPyStruct, WPyStructInfo>(s, info);
+    // steals a reference
+    PyList_SET_ITEM(a.ptr(), i, v.py.inc_ref().ptr());
+    ptr += sz;
+  }
+
+  return a;
+}
+
 // void unpackInto(const py::buffer &b, WPyStruct *v) {
 //   WPyStructInfo info(*v);
 //   py::ssize_t sz = wpi::GetStructSize<WPyStruct>(info);

subprojects/robotpy-wpiutil/wpiutil/src/wpistruct/wpystruct_fns.h

@@ -30,6 +30,11 @@ size_t getSize(const py::type &t);
 */
 py::bytes pack(const WPyStruct &v);
 
+/**
+    Serialize objects into byte buffer
+*/
+py::bytes packArray(const py::sequence &seq);
+
 /**
     Serialize object into byte buffer. Buffer must be exact size.
 */
@@ -41,6 +46,12 @@ void packInto(const WPyStruct &v, py::buffer &b);
 */
 WPyStruct unpack(const py::type &t, const py::buffer &b);
 
+/**
+    Convert byte buffer into list of objects of specified type. Buffer must be
+    exact size.
+*/
+py::typing::List<WPyStruct> unpackArray(const py::type &t, const py::buffer &b);
+
 // /**
 //     Convert byte buffer into passed in object. Buffer must be exact
 //     size.

subprojects/robotpy-wpiutil/wpiutil/wpistruct/__init__.py

@@ -10,8 +10,10 @@
     getSize,
     getTypeName,
     pack,
+    packArray,
     packInto,
     unpack,
+    unpackArray,
 )
 
 __all__ = [
@@ -20,8 +22,10 @@
     "getSize",
     "getTypeName",
     "pack",
+    "packArray",
     "packInto",
     "unpack",
+    "unpackArray",
 ]
 
 from .desc import StructDescriptor