From: Michael Tremer <michael.tremer@ipfire.org>
Date: Wed, 21 Apr 2021 11:10:36 +0000 (+0000)
Subject: progressbar: Implement using in Python context
X-Git-Tag: 0.9.28~1285^2~286
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=7a08b0f8ad053ec6439cca7ce8a0b1960f0a3708;p=pakfire.git

progressbar: Implement using in Python context

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

diff --git a/src/_pakfire/progressbar.c b/src/_pakfire/progressbar.c
index 5f4ca660f..8354df364 100644
--- a/src/_pakfire/progressbar.c
+++ b/src/_pakfire/progressbar.c
@@ -203,6 +203,28 @@ static PyObject* Progressbar_add_transfer_speed(ProgressbarObject* self) {
 	Py_RETURN_NONE;
 }
 
+static PyObject* Progressbar_enter(ProgressbarObject* self) {
+	int r = pakfire_progressbar_start(self->progressbar, 0);
+	if (r) {
+		PyErr_SetFromErrno(PyExc_OSError);
+		return NULL;
+	}
+
+	// Return self
+	Py_INCREF(self);
+	return (PyObject*)self;
+}
+
+static PyObject* Progressbar_exit(ProgressbarObject* self, PyObject* args) {
+	int r = pakfire_progressbar_finish(self->progressbar);
+	if (r) {
+		PyErr_SetFromErrno(PyExc_OSError);
+		return NULL;
+	}
+
+	Py_RETURN_NONE;
+}
+
 static struct PyMethodDef Progressbar_methods[] = {
 	{
 		"start",
@@ -288,6 +310,18 @@ static struct PyMethodDef Progressbar_methods[] = {
 		METH_NOARGS,
 		NULL,
 	},
+	{
+		"__enter__",
+		(PyCFunction)Progressbar_enter,
+		METH_NOARGS,
+		NULL,
+	},
+	{
+		"__exit__",
+		(PyCFunction)Progressbar_exit,
+		METH_VARARGS,
+		NULL,
+	},
 	{ NULL },
 };
 
diff --git a/tests/python/progressbar.py b/tests/python/progressbar.py
index 0c405738e..caf452ea8 100755
--- a/tests/python/progressbar.py
+++ b/tests/python/progressbar.py
@@ -33,6 +33,20 @@ class Test(unittest.TestCase):
 
 		p.finish()
 
+	def test_context(self):
+		p = Progressbar()
+
+		p.add_string("Progress...")
+		p.add_bar()
+
+		p.set_max(1000)
+
+		with p:
+			p.set_max(1000)
+
+			for i in range(1000):
+				p.increment()
+
 
 if __name__ == "__main__":
 	unittest.main()