python: execute: Allow returning the output of the command

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>

diff --git a/src/python/pakfire.c b/src/python/pakfire.c
index 6ccbda93b0dba3f9a0c3b9851960cbf55935d4f8..3268ac56127ba979b4050f29424f874ef5b61c79 100644 (file)
--- a/src/python/pakfire.c
+++ b/src/python/pakfire.c
@@ -532,6 +532,22 @@ ERROR:
  * Execute
  */
 
+static int Pakfire_execute_stdout_callback(
+               struct pakfire_ctx* ctx, void* data, const char* line, const size_t length) {
+       PyObject** output = data;
+       PyObject* chunk = NULL;
+
+       // Allocate a new chunk
+       chunk = PyBytes_FromStringAndSize(line, length);
+       if (!chunk)
+               return -ENOTSUP;
+
+       // Concatenate the chunk to the output
+       PyBytes_ConcatAndDel(output, chunk);
+
+       return 0;
+}
+
 static PyObject* Pakfire_execute(PakfireObject* self, PyObject* args, PyObject* kwargs) {
        struct pakfire_jail* jail = NULL;
        struct pakfire_env* env = NULL;
@@ -546,16 +562,22 @@ static PyObject* Pakfire_execute(PakfireObject* self, PyObject* args, PyObject*
        PyObject* environ = NULL;
        int nice = 0;
 
+       PyObject* output = NULL;
+       int return_output = 0;
+
+       pakfire_pty_stdout_callback stdout_callback = NULL;
+
        const char* kwlist[] = {
                "command",
                "environ",
                "nice",
+               "return_output",
                NULL,
        };
 
        // Parse arguments
-       if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|Oi", (char**)kwlist,
-                       &command, &environ, &nice))
+       if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|Oib", (char**)kwlist,
+                       &command, &environ, &nice, &return_output))
                goto ERROR;
 
        // Check if command is a list
@@ -624,6 +646,17 @@ static PyObject* Pakfire_execute(PakfireObject* self, PyObject* args, PyObject*
                }
        }
 
+       // Capture and return the output?
+       if (return_output) {
+               // Allocate a new buffer
+               output = PyBytes_FromString("");
+               if (!output)
+                       goto ERROR;
+
+               // Register a function that captures the output
+               stdout_callback = Pakfire_execute_stdout_callback;
+       }
+
        // Create a new jail
        r = pakfire_jail_create(&jail, self->pakfire);
        if (r < 0) {
@@ -645,7 +678,7 @@ static PyObject* Pakfire_execute(PakfireObject* self, PyObject* args, PyObject*
        Py_BEGIN_ALLOW_THREADS
 
        // Execute command
-       r = pakfire_jail_communicate(jail, argv, env, 0, NULL, NULL, NULL, NULL);
+       r = pakfire_jail_communicate(jail, argv, env, 0, NULL, NULL, stdout_callback, &output);
 
        Py_END_ALLOW_THREADS
 
@@ -666,8 +699,14 @@ static PyObject* Pakfire_execute(PakfireObject* self, PyObject* args, PyObject*
                goto ERROR;
        }
 
+       // Return the output if requested
+       if (return_output) {
+               ret = Py_NewRef(output);
+
        // Otherwise return None
-       ret = Py_NewRef(Py_None);
+       } else {
+               ret = Py_NewRef(Py_None);
+       }
 
 ERROR:
        if (jail)
@@ -677,6 +716,8 @@ ERROR:
        if (argv)
                free(argv);
 
+       Py_XDECREF(output);
+
        return ret;
 }
 

diff --git a/tests/python/execute.py b/tests/python/execute.py
index e2f8a2689d02e0c2d2e688905e541e1e6dd66dfc..3bed8e61792dea7dc5b6425cd7b32d690df9688f 100755 (executable)
--- a/tests/python/execute.py
+++ b/tests/python/execute.py
@@ -81,16 +81,23 @@ class ExecuteTests(tests.TestCase):
                        self.pakfire.execute(["/command-does-not-exist"])
 
        def test_execute_output(self):
-               self.pakfire.execute(["/command", "echo", "123"])
+               self.assertEqual(
+                       self.pakfire.execute(["/command", "echo", "123"], return_output=True),
+                       b"123\n",
+               )
 
                # Multiple newlines in one read
-               self.pakfire.execute(["/command", "echo", "1\n2\n3"])
+               self.assertEqual(
+                       self.pakfire.execute(["/command", "echo", "1\n2\n3"], return_output=True),
+                       b"1\n2\n3\n",
+               )
 
                # Run a command with a lot of output which exceeds the buffer size
                self.pakfire.execute(["/command", "lines", "1", "65536"])
 
                # Run a command that generates lots of lines
-               self.pakfire.execute(["/command", "lines", "100", "40"])
+               output = self.pakfire.execute(["/command", "lines", "100", "40"], return_output=True)
+               self.assertEquals(output.count(b"\n"), 100)
 
        def test_nice(self):
                self.pakfire.execute(["/command", "print-nice"], nice=5)