Make MapMutation a context manager

2 files changed, 27 insertions, 1 deletions

diff --git a/immutables/_map.c b/immutables/_map.c
index 1abf730..9f57582 100644
--- a/immutables/_map.c
+++ b/immutables/_map.c
@@ -3850,7 +3850,6 @@ mapmut_tp_richcompare(PyObject *v, PyObject *w, int op)
     }
 }
 
-
 static PyObject *
 mapmut_py_finalize(MapMutationObject *self, PyObject *args)
 {
@@ -3868,6 +3867,24 @@ mapmut_py_finalize(MapMutationObject *self, PyObject *args)
     return (PyObject *)o;
 }
 
+static PyObject *
+mapmut_py_enter(MapMutationObject *self, PyObject *args)
+{
+    Py_INCREF(self);
+    return (PyObject *)self;
+}
+
+static PyObject *
+mapmut_py_exit(MapMutationObject *self, PyObject *args)
+{
+    PyObject *ret = mapmut_py_finalize(self, NULL);
+    if (ret == NULL) {
+        return NULL;
+    }
+    Py_DECREF(ret);
+    Py_RETURN_FALSE;
+}
+
 static int
 mapmut_tp_ass_sub(MapMutationObject *self, PyObject *key, PyObject *val)
 {
@@ -3951,6 +3968,8 @@ static PyMethodDef MapMutation_methods[] = {
     {"get", (PyCFunction)map_py_get, METH_VARARGS, NULL},
     {"pop", (PyCFunction)mapmut_py_pop, METH_VARARGS, NULL},
     {"finalize", (PyCFunction)mapmut_py_finalize, METH_NOARGS, NULL},
+    {"__enter__", (PyCFunction)mapmut_py_enter, METH_NOARGS, NULL},
+    {"__exit__", (PyCFunction)mapmut_py_exit, METH_VARARGS, NULL},
     {NULL, NULL}
 };
 
diff --git a/immutables/map.py b/immutables/map.py
index abfe9ed..bea7ad9 100644
--- a/immutables/map.py
+++ b/immutables/map.py
@@ -633,6 +633,13 @@ class MapMutation:
     def set(self, key, val):
         self[key] = val
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        self.finalize()
+        return False
+
     def __delitem__(self, key):
         if self.__mutid == 0:
             raise ValueError(f'mutation {self!r} has been finalized')