From: Michael Tremer <michael.tremer@ipfire.org>
Date: Mon, 17 Mar 2025 11:46:29 +0000 (+0000)
Subject: python: execute: Allow returning the output of the command
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=608f9ea0b925f6b2110b3433a55f52e5a08c45d1;p=pakfire.git

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 6ccbda93..3268ac56 100644
--- 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 e2f8a268..3bed8e61 100755
--- 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)