python/ext/py-source.c: __exit__ to close once

[aubio.git] / python / ext / py-source.c

diff --git a/python/ext/py-source.c b/python/ext/py-source.c
index 55ed71e..82247fd 100644 (file)
--- a/python/ext/py-source.c
+++ b/python/ext/py-source.c
@@ -140,8 +140,8 @@ Py_source_init (Py_source * self, PyObject * args, PyObject * kwds)
 {
   self->o = new_aubio_source ( self->uri, self->samplerate, self->hop_size );
   if (self->o == NULL) {
-    PyErr_Format (PyExc_RuntimeError, "error creating source with \"%s\"",
-        self->uri);
+    // PyErr_Format(PyExc_RuntimeError, ...) was set above by new_ which called
+    // AUBIO_ERR when failing
     return -1;
   }
   self->samplerate = aubio_source_get_samplerate ( self->o );
@@ -173,6 +173,7 @@ Py_source_del (Py_source *self, PyObject *unused)
 static PyObject *
 Py_source_do(Py_source * self, PyObject * args)
 {
+  PyObject *outputs;
   uint_t read;
   read = 0;
 
@@ -183,7 +184,7 @@ Py_source_do(Py_source * self, PyObject * args)
   /* compute _do function */
   aubio_source_do (self->o, &(self->c_read_to), &read);
 
-  PyObject *outputs = PyTuple_New(2);
+  outputs = PyTuple_New(2);
   PyTuple_SetItem( outputs, 0, self->read_to );
   PyTuple_SetItem( outputs, 1, (PyObject *)PyLong_FromLong(read));
   return outputs;
@@ -193,6 +194,7 @@ Py_source_do(Py_source * self, PyObject * args)
 static PyObject *
 Py_source_do_multi(Py_source * self, PyObject * args)
 {
+  PyObject *outputs;
   uint_t read;
   read = 0;
 
@@ -203,7 +205,7 @@ Py_source_do_multi(Py_source * self, PyObject * args)
   /* compute _do function */
   aubio_source_do_multi (self->o, &(self->c_mread_to), &read);
 
-  PyObject *outputs = PyTuple_New(2);
+  outputs = PyTuple_New(2);
   PyTuple_SetItem( outputs, 0, self->mread_to);
   PyTuple_SetItem( outputs, 1, (PyObject *)PyLong_FromLong(read));
   return outputs;
@@ -240,7 +242,7 @@ Pyaubio_source_get_channels (Py_source *self, PyObject *unused)
 static PyObject *
 Pyaubio_source_close (Py_source *self, PyObject *unused)
 {
-  aubio_source_close (self->o);
+  if (aubio_source_close(self->o) != 0) return NULL;
   Py_RETURN_NONE;
 }
 
@@ -249,11 +251,18 @@ Pyaubio_source_seek (Py_source *self, PyObject *args)
 {
   uint_t err = 0;
 
-  uint_t position;
+  int position;
   if (!PyArg_ParseTuple (args, "I", &position)) {
     return NULL;
   }
 
+  if (position < 0) {
+    PyErr_Format(PyExc_ValueError,
+        "error when seeking in source: can not seek to negative value %d",
+        position);
+    return NULL;
+  }
+
   err = aubio_source_seek(self->o, position);
   if (err != 0) {
     PyErr_SetString (PyExc_ValueError,
@@ -263,6 +272,63 @@ Pyaubio_source_seek (Py_source *self, PyObject *args)
   Py_RETURN_NONE;
 }
 
+static char Pyaubio_source_enter_doc[] = "";
+static PyObject* Pyaubio_source_enter(Py_source *self, PyObject *unused) {
+  Py_INCREF(self);
+  return (PyObject*)self;
+}
+
+static char Pyaubio_source_exit_doc[] = "";
+static PyObject* Pyaubio_source_exit(Py_source *self, PyObject *unused) {
+  return Pyaubio_source_close(self, unused);
+}
+
+static PyObject* Pyaubio_source_iter(PyObject *self) {
+  Py_INCREF(self);
+  return (PyObject*)self;
+}
+
+static PyObject* Pyaubio_source_iter_next(Py_source *self) {
+  PyObject *done;
+  if (self->channels == 1) {
+    done = Py_source_do(self, NULL);
+  } else {
+    done = Py_source_do_multi(self, NULL);
+  }
+  if (!PyTuple_Check(done)) {
+    PyErr_Format(PyExc_ValueError,
+        "error when reading source: not opened?");
+    return NULL;
+  }
+  PyObject *size = PyTuple_GetItem(done, 1);
+  if (size != NULL && PyLong_Check(size)) {
+    if (PyLong_AsLong(size) == (long)self->hop_size) {
+      PyObject *vec = PyTuple_GetItem(done, 0);
+      return vec;
+    } else if (PyLong_AsLong(size) > 0) {
+      // short read
+      PyArrayObject *shortread = (PyArrayObject*)PyTuple_GetItem(done, 0);
+      PyArray_Dims newdims;
+      newdims.len = PyArray_NDIM(shortread);
+      newdims.ptr = PyArray_DIMS(shortread);
+      // mono or multiple channels?
+      if (newdims.len == 1) {
+        newdims.ptr[0] = PyLong_AsLong(size);
+      } else {
+        newdims.ptr[1] = PyLong_AsLong(size);
+      }
+      PyArray_Resize(shortread, &newdims, 1, NPY_ANYORDER);
+      return (PyObject*)shortread;
+    } else {
+      PyErr_SetNone(PyExc_StopIteration);
+      return NULL;
+    }
+  } else {
+    PyErr_SetNone(PyExc_StopIteration);
+    return NULL;
+  }
+}
+
 static PyMethodDef Py_source_methods[] = {
   {"get_samplerate", (PyCFunction) Pyaubio_source_get_samplerate,
     METH_NOARGS, Py_source_get_samplerate_doc},
@@ -276,6 +342,10 @@ static PyMethodDef Py_source_methods[] = {
     METH_NOARGS, Py_source_close_doc},
   {"seek", (PyCFunction) Pyaubio_source_seek,
     METH_VARARGS, Py_source_seek_doc},
+  {"__enter__", (PyCFunction)Pyaubio_source_enter, METH_NOARGS,
+    Pyaubio_source_enter_doc},
+  {"__exit__",  (PyCFunction)Pyaubio_source_exit, METH_VARARGS,
+    Pyaubio_source_exit_doc},
   {NULL} /* sentinel */
 };
 
@@ -305,8 +375,8 @@ PyTypeObject Py_sourceType = {
   0,
   0,
   0,
-  0,
-  0,
+  Pyaubio_source_iter,
+  (unaryfunc)Pyaubio_source_iter_next,
   Py_source_methods,
   Py_source_members,
   0,
@@ -318,4 +388,13 @@ PyTypeObject Py_sourceType = {
   (initproc) Py_source_init,
   0,
   Py_source_new,
+  0,
+  0,
+  0,
+  0,
+  0,
+  0,
+  0,
+  0,
+  0,
 };