catacomb.c (mexp_common): Accept an arbitrary iterable.

[catacomb-python] / catacomb.c

diff --git a/catacomb.c b/catacomb.c
index 1f38213..23906eb 100644 (file)
--- a/catacomb.c
+++ b/catacomb.c
@@ -38,61 +38,69 @@ PyObject *mexp_common(PyObject *me, PyObject *arg,
                      PyObject *(*exp)(PyObject *, void *, size_t),
                      void (*drop)(void *))
 {
-  size_t i = 0, j, n;
+  size_t i = 0, o, n;
   int flat;
-  PyObject *qq, *x, *y, *z = 0;
-  char *v = 0, *vv;
-
-  if (PyTuple_GET_SIZE(arg) == 1)
-    arg = PyTuple_GET_ITEM(arg, 0);
-  if (!PySequence_Check(arg)) TYERR("not a sequence");
-  n = PySequence_Size(arg); if (n < 0) goto end;
-  if (!n) { z = id(me); goto end; }
-  x = PySequence_GetItem(arg, 0);
-  if (PySequence_Check(x))
-    flat = 0;
+  PyObject *qq = 0, *x = 0, *y = 0, *z = 0, *it = 0;
+  char *v = 0;
+
+  if (PyTuple_Size(arg) == 1) arg = PyTuple_GET_ITEM(arg, 0);
+  it = PyObject_GetIter(arg); if (!it) goto end;
+  qq = PyIter_Next(it);
+  if (!qq) {
+    if (!PyErr_Occurred()) z = id(me);
+    else goto end;
+  }
+  flat = !PySequence_Check(qq);
+  if (!PySequence_Check(arg))
+    n = 16;
   else {
-    if (n % 2) VALERR("must have even number of arguments");
-    n /= 2;
-    flat = 1;
+    n = PySequence_Size(arg);
+    if (n == (size_t)-1 && PyErr_Occurred()) goto end;
+    if (flat) n /= 2;
+    if (!n) n = 16;
   }
-  Py_DECREF(x);
 
-  v = xmalloc(n * efsz);
-  for (i = j = 0, vv = v; i < n; i++, vv += efsz) {
-    if (flat) {
-      x = PySequence_GetItem(arg, j++);
-      y = PySequence_GetItem(arg, j++);
-    } else {
-      qq = PySequence_GetItem(arg, j++);
-      if (!qq) goto end;
-      if (!PySequence_Check(qq) || PySequence_Size(qq) != 2) {
-       Py_DECREF(qq);
+  v = xmalloc(n*efsz);
+  o = 0;
+  for (;;) {
+    if (!flat) {
+      if (!PySequence_Check(qq) || PySequence_Size(qq) != 2)
        TYERR("want a sequence of pairs");
-      }
       x = PySequence_GetItem(qq, 0);
       y = PySequence_GetItem(qq, 1);
-      Py_DECREF(qq);
+    } else {
+      x = qq; qq = 0;
+      y = PyIter_Next(it);
+      if (!y) {
+       if (PyErr_Occurred()) goto end;
+       VALERR("must have even number of operands");
+      }
     }
     if (!x || !y) goto end;
-    if (fill(vv, me, x, y)) {
-      Py_DECREF(x);
-      Py_DECREF(y);
-      if (!PyErr_Occurred())
-       PyErr_SetString(PyExc_TypeError, "type mismatch");
-      goto end;
+
+    if (i >= n) { n *= 2; v = xrealloc(v, n*efsz, i*efsz); }
+    if (fill(v + o, me, x, y)) {
+      if (PyErr_Occurred()) goto end;
+      TYERR("type mismatch");
+    }
+    i++; o += efsz;
+    Py_DECREF(x); x = 0;
+    Py_DECREF(y); y = 0;
+    Py_XDECREF(qq);
+
+    qq = PyIter_Next(it);
+    if (!qq) {
+      if (PyErr_Occurred()) goto end;
+      else break;
     }
-    Py_DECREF(x);
-    Py_DECREF(y);
   }
-  z = exp(me, v, n);
+
+  z = exp(me, v, i);
 
 end:
-  if (v) {
-    for (j = 0, vv = v; j < i; j++, vv += efsz)
-      drop(vv);
-    xfree(v);
-  }
+  while (i--) { o -= efsz; drop(v + o); }
+  xfree(v);
+  Py_XDECREF(it); Py_XDECREF(qq); Py_XDECREF(x); Py_XDECREF(y);
   return (z);
 }