diff --git a/include/rogue/interfaces/stream/Frame.h b/include/rogue/interfaces/stream/Frame.h index a1598a7d6..0189fe4b0 100644 --- a/include/rogue/interfaces/stream/Frame.h +++ b/include/rogue/interfaces/stream/Frame.h @@ -28,6 +28,8 @@ #include "rogue/EnableSharedFromThis.h" #ifndef NO_PYTHON + #include + #include #endif @@ -335,6 +337,22 @@ class Frame : public rogue::EnableSharedFromThis - #include + #include #include namespace bp = boost::python; @@ -346,7 +345,7 @@ ris::FrameIterator ris::Frame::endWrite() { #ifndef NO_PYTHON -//! Read up to count bytes from frame, starting from offset. Python version. +//! Read bytes from frame into a passed bytearray, starting from offset. Python version. void ris::Frame::readPy(boost::python::object p, uint32_t offset) { Py_buffer pyBuf; @@ -371,6 +370,44 @@ void ris::Frame::readPy(boost::python::object p, uint32_t offset) { PyBuffer_Release(&pyBuf); } +//! Allocate a bytearray and read bytes from frame into it starting at offset, return array +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, count))); + + // readPy will check bounds + this->readPy(byteArray, offset); + + return byteArray; +} + +//! Allocate a bytearray and from frame into it starting at offset, return memoryview to array +bp::object ris::Frame::getMemoryviewPy() { + // Get the size of the frame + uint32_t size = getPayload(); + + // Create a Python bytearray to hold the data + bp::object byteArray(bp::handle<>(PyByteArray_FromStringAndSize(nullptr, size))); + + this->readPy(byteArray, 0); + + // Create a memoryview from the bytearray + PyObject* memoryView = PyMemoryView_FromObject(byteArray.ptr()); + if (!memoryView) { + throw(rogue::GeneralError::create("Frame::getMemoryviewPy", "Failed to create memoryview.")); + } + + // Return the memoryview as a Python object + return bp::object(bp::handle<>(memoryView)); +} + //! Write python buffer to frame, starting at offset. Python Version void ris::Frame::writePy(boost::python::object p, uint32_t offset) { Py_buffer pyBuf; @@ -398,10 +435,14 @@ void ris::Frame::writePy(boost::python::object p, uint32_t offset) { } //! Read the specified number of bytes at the specified offset of frame data into a numpy array -boost::python::object ris::Frame::getNumpy(uint32_t offset, uint32_t count) { +boost::python::object ris::Frame::getNumpy(uint32_t offset, uint32_t count, bp::object dtype) { // Retrieve the size, in bytes of the data npy_intp size = getPayload(); + if (count == 0) { + count = size - offset; + } + // Check this does not request data past the EOF if ((offset + count) > size) { throw(rogue::GeneralError::create("Frame::getNumpy", @@ -412,9 +453,18 @@ boost::python::object ris::Frame::getNumpy(uint32_t offset, uint32_t count) { size)); } + // Convert Python dtype object to NumPy type + int numpy_type; + PyObject* dtype_pyobj = dtype.ptr(); // Get the raw PyObject from the Boost.Python object + if (PyArray_DescrCheck(dtype_pyobj)) { + numpy_type = (reinterpret_cast(dtype_pyobj))->type_num; + } else { + throw(rogue::GeneralError::create("Frame::getNumpy", "Invalid dtype argument. Must be a NumPy dtype object.")); + } + // Create a numpy array to receive it and locate the destination data buffer npy_intp dims[1] = {count}; - PyObject* obj = PyArray_SimpleNew(1, dims, NPY_UINT8); + PyObject* obj = PyArray_SimpleNew(1, dims, numpy_type); PyArrayObject* arr = reinterpret_cast(obj); uint8_t* dst = reinterpret_cast(PyArray_DATA(arr)); @@ -486,13 +536,22 @@ void ris::Frame::setup_python() { _import_array(); + // Create a NumPy dtype object from the NPY_UINT8 constant + PyObject* dtype_uint8 = reinterpret_cast(PyArray_DescrFromType(NPY_UINT8)); + if (!dtype_uint8) { + throw( + rogue::GeneralError::create("Frame::setup_python", "Failed to create default dtype object for NPY_UINT8.")); + } + bp::class_("Frame", bp::no_init) .def("lock", &ris::Frame::lock) .def("getSize", &ris::Frame::getSize) .def("getAvailable", &ris::Frame::getAvailable) .def("getPayload", &ris::Frame::getPayload) - .def("read", &ris::Frame::readPy) - .def("write", &ris::Frame::writePy) + .def("read", &ris::Frame::readPy, (bp::arg("offset") = 0)) + .def("getBa", &ris::Frame::getBytearrayPy, (bp::arg("offset") = 0, bp::arg("count") = 0)) + .def("getMemoryview", &ris::Frame::getMemoryviewPy) + .def("write", &ris::Frame::writePy, (bp::arg("offset") = 0)) .def("setError", &ris::Frame::setError) .def("getError", &ris::Frame::getError) .def("setFlags", &ris::Frame::setFlags) @@ -503,8 +562,12 @@ void ris::Frame::setup_python() { .def("getLastUser", &ris::Frame::getLastUser) .def("setChannel", &ris::Frame::setChannel) .def("getChannel", &ris::Frame::getChannel) - .def("getNumpy", &ris::Frame::getNumpy) - .def("putNumpy", &ris::Frame::putNumpy) + .def("getNumpy", + &ris::Frame::getNumpy, + (bp::arg("offset") = 0, + bp::arg("count") = 0, + bp::arg("dtype") = bp::object(bp::handle<>(bp::borrowed(dtype_uint8))))) + .def("putNumpy", &ris::Frame::putNumpy, (bp::arg("offset") = 0)) .def("_debug", &ris::Frame::debug); #endif }