[pygobject] Add context management to freeze_notify() and handler_block().

[Martin Pitt <martinpitt src gnome org>]

commit 0fd900d351c8d7d57dc6a1b049ee05f342f6ab1d
Author: Simon Feltman <s feltman gmail com>
Date:   Sun Mar 18 15:59:58 2012 -0700

    Add context management to freeze_notify() and handler_block().
    
    These methods now return a context manager object. Within the __exit__ method
    thaw_notify() and handler_unblock() are called respectively. This allows
    statements like the following:
    
    with obj.freeze_notify():
        obj.props.width = 100
        obj.props.height = 100
        obj.props.opacity = 0.5
    
    This does not affect standard usage of these methods.
    
    https://bugzilla.gnome.org/show_bug.cgi?id=672324
    
    Signed-off-by: Martin Pitt <martinpitt gnome org>

 gi/_gobject/pygobject.c |  138 ++++++++++++++++++++++++++++++++++++++---
 tests/test_gobject.py   |  158 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 288 insertions(+), 8 deletions(-)
---
diff --git a/gi/_gobject/pygobject.c b/gi/_gobject/pygobject.c
index 66b8ae6..9d20b97 100644
--- a/gi/_gobject/pygobject.c
+++ b/gi/_gobject/pygobject.c
@@ -993,6 +993,116 @@ pygobject_watch_closure(PyObject *self, GClosure *closure)
     g_closure_add_invalidate_notifier(closure, data, pygobject_unwatch_closure);
 }
 
+
+/* -------------- Freeze Notify Context Manager ----------------- */
+
+/**
+ * pygcontext_manager_enter
+ * @self: Freeze or Block context instance
+ *
+ * Method used for __enter__ on both GContextFeezeNotify and
+ * GContextHandlerBlock. Does nothing since this is an object returned
+ * by the freeze_notify() and handler_block() methods which do the actual
+ * work of freezing and blocking.
+ */
+static PyObject *
+pygcontext_manager_enter(PyObject *self)
+{
+	Py_INCREF(self);
+	return self;
+}
+
+typedef struct {
+    PyObject_HEAD
+    GObject *obj;
+} PyGContextFreezeNotify;
+
+PYGLIB_DEFINE_TYPE("gi._gobject.GContextFreezeNotify",
+		PyGContextFreezeNotify_Type, PyGContextFreezeNotify);
+
+static PyObject *
+pygcontext_freeze_notify_new(GObject *gobj)
+{
+	PyGContextFreezeNotify *context;
+
+	context = PyObject_New(PyGContextFreezeNotify, &PyGContextFreezeNotify_Type);
+	if (context == NULL)
+		return NULL;
+
+	g_object_ref(gobj);
+	context->obj = gobj;
+	return (PyObject*)context;
+}
+
+static void
+pygcontext_freeze_notify_dealloc(PyGContextFreezeNotify* self)
+{
+    g_object_unref(self->obj);
+    self->obj = NULL;
+    PyObject_Del((PyObject*)self);
+}
+
+static PyObject *
+pygcontext_freeze_notify_exit(PyGContextFreezeNotify *self, PyObject *args)
+{
+	g_object_thaw_notify(self->obj);
+	Py_RETURN_NONE;
+}
+
+static PyMethodDef pygcontext_freeze_notify_methods[] = {
+	{"__enter__", (PyCFunction)pygcontext_manager_enter, METH_NOARGS, ""},
+	{"__exit__", (PyCFunction)pygcontext_freeze_notify_exit, METH_VARARGS, ""},
+    {NULL}
+};
+
+/* -------------- Handler Block Context Manager ----------------- */
+typedef struct {
+    PyObject_HEAD
+    GObject *obj;
+    gulong handler_id;
+} PyGContextHandlerBlock;
+
+PYGLIB_DEFINE_TYPE("gi._gobject.GContextHandlerBlock",
+		PyGContextHandlerBlock_Type, PyGContextHandlerBlock);
+
+static PyObject *
+pygcontext_handler_block_new(GObject *gobj, gulong handler_id)
+{
+	PyGContextHandlerBlock *context;
+
+	context = PyObject_New(PyGContextHandlerBlock, &PyGContextHandlerBlock_Type);
+	if (context == NULL)
+		return NULL;
+
+	g_object_ref(gobj);
+	context->obj = gobj;
+	context->handler_id = handler_id;
+	return (PyObject*)context;
+}
+
+static void
+pygcontext_handler_block_dealloc(PyGContextHandlerBlock* self)
+{
+    g_object_unref(self->obj);
+    self->obj = NULL;
+    PyObject_Del((PyObject*)self);
+}
+
+static PyObject *
+pygcontext_handler_block_exit(PyGContextHandlerBlock *self, PyObject *args)
+{
+	g_signal_handler_unblock(self->obj, self->handler_id);
+	Py_RETURN_NONE;
+}
+
+
+static PyMethodDef pygcontext_handler_block_methods[] = {
+	{"__enter__", (PyCFunction)pygcontext_manager_enter, METH_NOARGS, ""},
+	{"__exit__", (PyCFunction)pygcontext_handler_block_exit, METH_VARARGS, ""},
+    {NULL}
+};
+
+
 /* -------------- PyGObject behaviour ----------------- */
 
 PYGLIB_DEFINE_TYPE("gi._gobject.GObject", PyGObject_Type, PyGObject);
@@ -1380,12 +1490,11 @@ pygobject_freeze_notify(PyGObject *self, PyObject *args)
 {
     if (!PyArg_ParseTuple(args, ":GObject.freeze_notify"))
 	return NULL;
-    
+
     CHECK_GOBJECT(self);
-    
+
     g_object_freeze_notify(self->obj);
-    Py_INCREF(Py_None);
-    return Py_None;
+    return pygcontext_freeze_notify_new(self->obj);
 }
 
 static PyObject *
@@ -1395,9 +1504,9 @@ pygobject_notify(PyGObject *self, PyObject *args)
 
     if (!PyArg_ParseTuple(args, "s:GObject.notify", &property_name))
 	return NULL;
-    
+
     CHECK_GOBJECT(self);
-    
+
     g_object_notify(self->obj, property_name);
     Py_INCREF(Py_None);
     return Py_None;
@@ -1666,8 +1775,7 @@ pygobject_handler_block(PyGObject *self, PyObject *args)
     CHECK_GOBJECT(self);
     
     g_signal_handler_block(self->obj, handler_id);
-    Py_INCREF(Py_None);
-    return Py_None;
+    return pygcontext_handler_block_new(self->obj, handler_id);
 }
 
 static PyObject *
@@ -2318,4 +2426,18 @@ pygobject_object_register_types(PyObject *d)
     if (PyType_Ready(&PyGObjectWeakRef_Type) < 0)
         return;
     PyDict_SetItemString(d, "GObjectWeakRef", (PyObject *) &PyGObjectWeakRef_Type);
+
+    PyGContextFreezeNotify_Type.tp_dealloc = (destructor)pygcontext_freeze_notify_dealloc;
+    PyGContextFreezeNotify_Type.tp_flags = Py_TPFLAGS_DEFAULT;
+    PyGContextFreezeNotify_Type.tp_doc = "Context manager for freeze/thaw of GObjects";
+    PyGContextFreezeNotify_Type.tp_methods = pygcontext_freeze_notify_methods;
+    if (PyType_Ready(&PyGContextFreezeNotify_Type) < 0)
+        return;
+
+    PyGContextHandlerBlock_Type.tp_dealloc = (destructor)pygcontext_handler_block_dealloc;
+    PyGContextHandlerBlock_Type.tp_flags = Py_TPFLAGS_DEFAULT;
+    PyGContextHandlerBlock_Type.tp_doc = "Context manager for handler blocking of GObjects";
+    PyGContextHandlerBlock_Type.tp_methods = pygcontext_handler_block_methods;
+    if (PyType_Ready(&PyGContextHandlerBlock_Type) < 0)
+        return;
 }
diff --git a/tests/test_gobject.py b/tests/test_gobject.py
index 80725b3..ecd67cf 100644
--- a/tests/test_gobject.py
+++ b/tests/test_gobject.py
@@ -190,3 +190,161 @@ class TestPythonReferenceCounting(unittest.TestCase):
     def testNewSubclassInstanceHasTwoRefsUsingGObjectNew(self):
         obj = GObject.new(A)
         self.assertEquals(sys.getrefcount(obj), 2)
+
+
+class TestContextManagers(unittest.TestCase):
+    class ContextTestObject(GObject.GObject):
+        prop = GObject.Property(default=0, type=int)
+
+    def on_prop_set(self, obj, prop):
+        # Handler which tracks property changed notifications.
+        self.tracking.append(obj.get_property(prop.name))
+
+    def setUp(self):
+        self.tracking = []
+        self.obj = self.ContextTestObject()
+        self.handler = self.obj.connect('notify::prop', self.on_prop_set)
+
+    def testFreezeNotifyContext(self):
+        # Verify prop tracking list
+        self.assertEqual(self.tracking, [])
+        self.obj.props.prop = 1
+        self.assertEqual(self.tracking, [1])
+        self.obj.props.prop = 2
+        self.assertEqual(self.tracking, [1, 2])
+        self.assertEqual(self.obj.__grefcount__, 1)
+
+        # Using the context manager the tracking list should not be affected
+        # and the GObject reference count should go up.
+        with self.obj.freeze_notify():
+            self.assertEqual(self.obj.__grefcount__, 2)
+            self.obj.props.prop = 3
+            self.assertEqual(self.obj.props.prop, 3)
+            self.assertEqual(self.tracking, [1, 2])
+
+        # After the context manager, the prop should have been modified,
+        # the tracking list will be modified, and the GObject ref
+        # count goes back down.
+        self.assertEqual(self.obj.props.prop, 3)
+        self.assertEqual(self.tracking, [1, 2, 3])
+        self.assertEqual(self.obj.__grefcount__, 1)
+
+    def testHandlerBlockContext(self):
+        # Verify prop tracking list
+        self.assertEqual(self.tracking, [])
+        self.obj.props.prop = 1
+        self.assertEqual(self.tracking, [1])
+        self.obj.props.prop = 2
+        self.assertEqual(self.tracking, [1, 2])
+        self.assertEqual(self.obj.__grefcount__, 1)
+
+        # Using the context manager the tracking list should not be affected
+        # and the GObject reference count should go up.
+        with self.obj.handler_block(self.handler):
+            self.assertEqual(self.obj.__grefcount__, 2)
+            self.obj.props.prop = 3
+            self.assertEqual(self.obj.props.prop, 3)
+            self.assertEqual(self.tracking, [1, 2])
+
+        # After the context manager, the prop should have been modified
+        # the tracking list should have stayed the same and the GObject ref
+        # count goes back down.
+        self.assertEqual(self.obj.props.prop, 3)
+        self.assertEqual(self.tracking, [1, 2])
+        self.assertEqual(self.obj.__grefcount__, 1)
+
+    def testFreezeNotifyContextNested(self):
+        self.assertEqual(self.tracking, [])
+        with self.obj.freeze_notify():
+            self.obj.props.prop = 1
+            self.assertEqual(self.tracking, [])
+
+            with self.obj.freeze_notify():
+                self.obj.props.prop = 2
+                self.assertEqual(self.tracking, [])
+
+                with self.obj.freeze_notify():
+                    self.obj.props.prop = 3
+                    self.assertEqual(self.tracking, [])
+                self.assertEqual(self.tracking, [])
+            self.assertEqual(self.tracking, [])
+
+        # Finally after last context, the notifications should have collapsed
+        # and the last one sent.
+        self.assertEqual(self.tracking, [3])
+
+    def testHandlerBlockContextNested(self):
+        self.assertEqual(self.tracking, [])
+        with self.obj.handler_block(self.handler):
+            self.obj.props.prop = 1
+            self.assertEqual(self.tracking, [])
+
+            with self.obj.handler_block(self.handler):
+                self.obj.props.prop = 2
+                self.assertEqual(self.tracking, [])
+
+                with self.obj.handler_block(self.handler):
+                    self.obj.props.prop = 3
+                    self.assertEqual(self.tracking, [])
+                self.assertEqual(self.tracking, [])
+            self.assertEqual(self.tracking, [])
+
+        # Finally after last context, the notifications should have collapsed
+        # and the last one sent.
+        self.assertEqual(self.obj.props.prop, 3)
+        self.assertEqual(self.tracking, [])
+
+    def testFreezeNotifyNormalUsageRefCounts(self):
+        # Ensure ref counts without using methods as context managers
+        # maintain the same count.
+        self.assertEqual(self.obj.__grefcount__, 1)
+        self.obj.freeze_notify()
+        self.assertEqual(self.obj.__grefcount__, 1)
+        self.obj.thaw_notify()
+        self.assertEqual(self.obj.__grefcount__, 1)
+
+    def testHandlerBlockNormalUsageRefCounts(self):
+        self.assertEqual(self.obj.__grefcount__, 1)
+        self.obj.handler_block(self.handler)
+        self.assertEqual(self.obj.__grefcount__, 1)
+        self.obj.handler_unblock(self.handler)
+        self.assertEqual(self.obj.__grefcount__, 1)
+
+    def testFreezeNotifyContextError(self):
+        # Test an exception occurring within a freeze context exits the context
+        try:
+            with self.obj.freeze_notify():
+                self.obj.props.prop = 1
+                self.assertEqual(self.tracking, [])
+                raise ValueError('Simulation')
+        except ValueError:
+            pass
+
+        # Verify the property set within the context called notify.
+        self.assertEqual(self.obj.props.prop, 1)
+        self.assertEqual(self.tracking, [1])
+
+        # Verify we are still not in a frozen context.
+        self.obj.props.prop = 2
+        self.assertEqual(self.tracking, [1, 2])
+
+    def testHandlerBlockContextError(self):
+        # Test an exception occurring within a handler block exits the context
+        try:
+            with self.obj.handler_block(self.handler):
+                self.obj.props.prop = 1
+                self.assertEqual(self.tracking, [])
+                raise ValueError('Simulation')
+        except ValueError:
+            pass
+
+        # Verify the property set within the context didn't call notify.
+        self.assertEqual(self.obj.props.prop, 1)
+        self.assertEqual(self.tracking, [])
+
+        # Verify we are still not in a handler block context.
+        self.obj.props.prop = 2
+        self.assertEqual(self.tracking, [2])
+
+if __name__ == '__main__':
+    unittest.main()