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 5f4ca660f8030e437cd799d3eb93527f0f697b33..8354df3645d1eeb84b65c68500f36fd726412249 100644 (file)
--- 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 0c405738e122f6389960d643b2c955b062dfd613..caf452ea88b57c2ad6f163f07aa2a4dd9ed62085 100755 (executable)
--- 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()