Add count parameter to getBytearrayPy

slaclab · Oct 16, 2024 · a2f3067 · a2f3067
1 parent 86a87ed
commit a2f3067
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/include/rogue/interfaces/stream/Frame.h b/include/rogue/interfaces/stream/Frame.h
@@ -343,8 +343,9 @@ class Frame : public rogue::EnableSharedFromThis<rogue::interfaces::stream::Fram
      *
      * Exposed as getBa() to Python
      * @param offset First location of Frame data to copy to byte array
+     * @param count Number of bytes to read
      */
-    boost::python::object getBytearrayPy(uint32_t offset);
+    boost::python::object getBytearrayPy(uint32_t offset, uint32_t count);
 
     //! Python Frame data read function
     /** Read data from Frame into a python bytearray which is allocated and returned as a memoryview

diff --git a/src/rogue/interfaces/stream/Frame.cpp b/src/rogue/interfaces/stream/Frame.cpp
@@ -371,14 +371,19 @@ void ris::Frame::readPy(boost::python::object p, uint32_t offset) {
 }
 
 //! Allocate a bytearray and read bytes from frame into it starting at offset, return array
-bp::object ris::Frame::getBytearrayPy(uint32_t offset) {
+bp::object ris::Frame::getBytearrayPy(uint32_t offset, uint32_t count) {
     // Get the size of the frame
     uint32_t size = getPayload();
+
+    if (count == 0) {
+        count = size - offset;
+    }
+
 
     // Create a Python bytearray to hold the data
-    bp::object byteArray(bp::handle<>(PyByteArray_FromStringAndSize(nullptr, size - offset)));
-
+    bp::object byteArray(bp::handle<>(PyByteArray_FromStringAndSize(nullptr, count)));
 
+    // readPy will check bounds
     this->readPy(byteArray, offset);
 
     return byteArray;
@@ -554,7 +559,8 @@ void ris::Frame::setup_python() {
         .def("read", &ris::Frame::readPy, (
             bp::arg("offset")=0))
         .def("getBa", &ris::Frame::getBytearrayPy, (
-            bp::arg("offset")=0))
+            bp::arg("offset")=0,
+            bp::arg("count")=0))
         .def("getMemoryview", &ris::Frame::getMemoryviewPy)
         .def("write", &ris::Frame::writePy, (
             bp::arg("offset")=0))