Bindings for vstack atom

sparsediffpy/_bindings/atoms/vstack.h

@@ -0,0 +1,50 @@
+#ifndef ATOM_VSTACK_H
+#define ATOM_VSTACK_H
+
+#include "common.h"
+
+static PyObject *py_make_vstack(PyObject *self, PyObject *args)
+{
+    PyObject *list_obj;
+    if (!PyArg_ParseTuple(args, "O", &list_obj))
+    {
+        return NULL;
+    }
+    if (!PyList_Check(list_obj))
+    {
+        PyErr_SetString(PyExc_TypeError,
+                        "First argument must be a list of expr capsules");
+        return NULL;
+    }
+    Py_ssize_t n_args = PyList_Size(list_obj);
+    if (n_args == 0)
+    {
+        PyErr_SetString(PyExc_ValueError, "List of expr capsules cannot be empty");
+        return NULL;
+    }
+    expr **expr_args = (expr **) calloc(n_args, sizeof(expr *));
+    for (Py_ssize_t i = 0; i < n_args; ++i)
+    {
+        PyObject *item = PyList_GetItem(list_obj, i);
+        expr *e = (expr *) PyCapsule_GetPointer(item, EXPR_CAPSULE_NAME);
+        if (!e)
+        {
+            free(expr_args);
+            PyErr_SetString(PyExc_ValueError, "Invalid expr capsule in list");
+            return NULL;
+        }
+        expr_args[i] = e;
+    }
+    int n_vars = expr_args[0]->n_vars;
+    expr *node = new_vstack(expr_args, (int) n_args, n_vars);
+    free(expr_args);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create vstack node");
+        return NULL;
+    }
+    expr_retain(node);
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif // ATOM_VSTACK_H

sparsediffpy/_bindings/bindings.c

@@ -46,6 +46,7 @@
 #include "atoms/trace.h"
 #include "atoms/transpose.h"
 #include "atoms/variable.h"
+#include "atoms/vstack.h"
 #include "atoms/xexp.h"
 
 /* Include problem bindings */
@@ -82,6 +83,9 @@ static PyMethodDef DNLPMethods[] = {
     {"make_hstack", py_make_hstack, METH_VARARGS,
      "Create hstack node from list of expr capsules and n_vars (make_hstack([e1, "
      "e2, ...], n_vars))"},
+    {"make_vstack", py_make_vstack, METH_VARARGS,
+     "Create vstack node from list of expr capsules (make_vstack([e1, e2, "
+     "...]))"},
     {"make_sum", py_make_sum, METH_VARARGS, "Create sum node"},
     {"make_neg", py_make_neg, METH_VARARGS, "Create neg node"},
     {"make_normal_cdf", py_make_normal_cdf, METH_VARARGS, "Create normal_cdf node"},