From c214680432cc5b101760de821592593fc1915d72 Mon Sep 17 00:00:00 2001
From: Andrew Svetlov <andrew.svetlov@gmail.com>
Date: Fri, 3 Jan 2025 21:00:01 +0100
Subject: [PATCH] Implement context manager for Token

---
 Lib/test/test_context.py | 25 +++++++++++++++++++++++++
 Python/context.c         | 20 ++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/Lib/test/test_context.py b/Lib/test/test_context.py
index 82d1797ab3b79e..33765f45124f9f 100644
--- a/Lib/test/test_context.py
+++ b/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
 
diff --git a/Python/context.c b/Python/context.c
index 95aa82206270f9..e98db9f99b3d77 100644
--- a/Python/context.c
+++ b/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}
 };