Implement context manager for Token

Author: asvetlov

Lib/test/test_context.py

@@ -383,6 +383,31 @@ def sub(num):
             tp.shutdown()
         self.assertEqual(results, list(range(10)))
 
+    def test_token_contextmanager_with_default(self):
+        ctx = contextvars.Context()
+        c = contextvars.ContextVar('c', default=42)
+
+        def fun():
+            with c.set(36):
+                self.assertEqual(c.get(), 36)
+
+            self.assertEqual(c.get(), 42)
+
+        ctx.run(fun)
+
+    def test_token_contextmanager_without_default(self):
+        ctx = contextvars.Context()
+        c = contextvars.ContextVar('c')
+
+        def fun():
+            with c.set(36):
+                self.assertEqual(c.get(), 36)
+
+            with self.assertRaisesRegex(LookupError, "<ContextVar name='c'"):
+                c.get()
+
+        ctx.run(fun)
+
 
 # HAMT Tests
 

Python/context.c

@@ -1216,9 +1216,29 @@ static PyGetSetDef PyContextTokenType_getsetlist[] = {
     {NULL}
 };
 
+static PyObject *
+token_enter(PyContextToken *self, PyObject *args)
+{
+    return Py_NewRef(self);
+}
+
+static PyObject *
+token_exit(PyContextToken *self, PyObject *args)
+{
+    int ret = PyContextVar_Reset((PyObject *)self->tok_var, (PyObject *)self);
+    if (ret < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef PyContextTokenType_methods[] = {
     {"__class_getitem__",    Py_GenericAlias,
     METH_O|METH_CLASS,       PyDoc_STR("See PEP 585")},
+
+    {"__enter__", (PyCFunction)token_enter, METH_NOARGS, NULL},
+    {"__exit__", (PyCFunction)token_exit, METH_VARARGS, NULL},
+
     {NULL}
 };