python: Wrap Analog::data() as a NumPy array.

[libsigrok.git] / bindings / python / sigrok / core / classes.i

diff --git a/bindings/python/sigrok/core/classes.i b/bindings/python/sigrok/core/classes.i
index b05f21e1be0284ea5454b8fbac7cf9020557c304..7d4304fca13e9644a2117d74571576fd24cd2495 100644 (file)
--- a/bindings/python/sigrok/core/classes.i
+++ b/bindings/python/sigrok/core/classes.i
@@ -46,6 +46,7 @@ which provides access to the error code and description."
 
 %{
 #include <pygobject.h>
+#include <numpy/arrayobject.h>
 
 PyObject *GLib;
 PyTypeObject *IOChannel;
@@ -66,6 +67,7 @@ typedef guint pyg_flags_type;
     GLib = PyImport_ImportModule("gi.repository.GLib");
     IOChannel = (PyTypeObject *) PyObject_GetAttrString(GLib, "IOChannel");
     PollFD = (PyTypeObject *) PyObject_GetAttrString(GLib, "PollFD");
+    import_array();
 %}
 
 /* Map file objects to file descriptors. */
@@ -173,15 +175,29 @@ typedef guint pyg_flags_type;
         Py_XDECREF(arglist);
         Py_XDECREF(revents_obj);
 
-        if (PyErr_Occurred() || !PyBool_Check(result))
-            throw sigrok::Error(SR_ERR);
+        bool completed = !PyErr_Occurred();
+
+        if (!completed)
+            PyErr_Print();
 
-        bool retval = (result == Py_True);
+        bool valid_result = (completed && PyBool_Check(result));
+
+        if (completed && !valid_result)
+        {
+            PyErr_SetString(PyExc_TypeError,
+                "EventSource callback did not return a boolean");
+            PyErr_Print();
+        }
+
+        bool retval = (valid_result && result == Py_True);
 
         Py_XDECREF(result);
 
         PyGILState_Release(gstate);
 
+        if (!valid_result)
+            throw sigrok::Error(SR_ERR);
+
         return retval;
     };
 
@@ -212,9 +228,27 @@ typedef guint pyg_flags_type;
         Py_XDECREF(arglist);
         Py_XDECREF(log_obj);
         Py_XDECREF(string_obj);
+
+        bool completed = !PyErr_Occurred();
+
+        if (!completed)
+            PyErr_Print();
+
+        bool valid_result = (completed && result == Py_None);
+
         Py_XDECREF(result);
 
+        if (completed && !valid_result)
+        {
+            PyErr_SetString(PyExc_TypeError,
+                "Log callback did not return None");
+            PyErr_Print();
+        }
+
         PyGILState_Release(gstate);
+
+        if (!valid_result)
+            throw sigrok::Error(SR_ERR);
     };
 
     Py_XINCREF($input);
@@ -248,14 +282,75 @@ typedef guint pyg_flags_type;
         Py_XDECREF(arglist);
         Py_XDECREF(device_obj);
         Py_XDECREF(packet_obj);
+
+        bool completed = !PyErr_Occurred();
+
+        if (!completed)
+            PyErr_Print();
+
+        bool valid_result = (completed && result == Py_None);
+
         Py_XDECREF(result);
 
+        if (completed && !valid_result)
+        {
+            PyErr_SetString(PyExc_TypeError,
+                "Datafeed callback did not return None");
+            PyErr_Print();
+        }
+
         PyGILState_Release(gstate);
+
+        if (!valid_result)
+            throw sigrok::Error(SR_ERR);
     };
 
     Py_XINCREF($input);
 }
 
+/* Cast PacketPayload pointers to correct subclass type. */
+%ignore sigrok::Packet::payload;
+
+%extend sigrok::Packet
+{
+    std::shared_ptr<sigrok::Header> _payload_header()
+    {
+        return dynamic_pointer_cast<sigrok::Header>($self->payload());
+    }
+    std::shared_ptr<sigrok::Meta> _payload_meta()
+    {
+        return dynamic_pointer_cast<sigrok::Meta>($self->payload());
+    }
+    std::shared_ptr<sigrok::Analog> _payload_analog()
+    {
+        return dynamic_pointer_cast<sigrok::Analog>($self->payload());
+    }
+    std::shared_ptr<sigrok::Logic> _payload_logic()
+    {
+        return dynamic_pointer_cast<sigrok::Logic>($self->payload());
+    }
+}
+
+%extend sigrok::Packet
+{
+%pythoncode
+{
+    def _payload(self):
+        if self.type == PacketType.HEADER:
+            return self._payload_header()
+        elif self.type == PacketType.META:
+            return self._payload_meta()
+        elif self.type == PacketType.LOGIC:
+            return self._payload_logic()
+        elif self.type == PacketType.ANALOG:
+            return self._payload_analog()
+        else:
+            return None
+
+    payload = property(_payload)
+}
+}
+
 %{
 
 #include "libsigrok/libsigrok.hpp"
@@ -353,6 +448,7 @@ std::map<std::string, Glib::VariantBase> dict_to_map_options(PyObject *dict,
 %}
 
 /* Ignore these methods, we will override them below. */
+%ignore sigrok::Analog::data;
 %ignore sigrok::Driver::scan;
 %ignore sigrok::InputFormat::create_input;
 %ignore sigrok::OutputFormat::create_output;
@@ -480,3 +576,23 @@ std::map<std::string, Glib::VariantBase> dict_to_map_options(PyObject *dict,
         $self->config_set(key, python_to_variant_by_key(input, key));
     }
 }
+
+/* Return NumPy array from Analog::data(). */
+%extend sigrok::Analog
+{
+    PyObject * _data()
+    {
+        int nd = 2;
+        npy_intp dims[2];
+        dims[0] = $self->channels().size();
+        dims[1] = $self->num_samples();
+        int typenum = NPY_FLOAT;
+        void *data = $self->data_pointer();
+        return PyArray_SimpleNewFromData(nd, dims, typenum, data);
+    }
+
+%pythoncode
+{
+    data = property(_data)
+}
+}