gh-129889: Support context manager protocol by contextvars.Token (#129888)

2 files changed, 90 insertions, 1 deletions

diff --git a/Python/clinic/context.c.h b/Python/clinic/context.c.h
index 71f05aa02a5..0adde76d7c3 100644
--- a/Python/clinic/context.c.h
+++ b/Python/clinic/context.c.h
@@ -179,4 +179,55 @@ PyDoc_STRVAR(_contextvars_ContextVar_reset__doc__,
 
 #define _CONTEXTVARS_CONTEXTVAR_RESET_METHODDEF    \
     {"reset", (PyCFunction)_contextvars_ContextVar_reset, METH_O, _contextvars_ContextVar_reset__doc__},
-/*[clinic end generated code: output=444567eaf0df25e0 input=a9049054013a1b77]*/
+
+PyDoc_STRVAR(token_enter__doc__,
+"__enter__($self, /)\n"
+"--\n"
+"\n"
+"Enter into Token context manager.");
+
+#define TOKEN_ENTER_METHODDEF    \
+    {"__enter__", (PyCFunction)token_enter, METH_NOARGS, token_enter__doc__},
+
+static PyObject *
+token_enter_impl(PyContextToken *self);
+
+static PyObject *
+token_enter(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    return token_enter_impl((PyContextToken *)self);
+}
+
+PyDoc_STRVAR(token_exit__doc__,
+"__exit__($self, type, val, tb, /)\n"
+"--\n"
+"\n"
+"Exit from Token context manager, restore the linked ContextVar.");
+
+#define TOKEN_EXIT_METHODDEF    \
+    {"__exit__", _PyCFunction_CAST(token_exit), METH_FASTCALL, token_exit__doc__},
+
+static PyObject *
+token_exit_impl(PyContextToken *self, PyObject *type, PyObject *val,
+                PyObject *tb);
+
+static PyObject *
+token_exit(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
+{
+    PyObject *return_value = NULL;
+    PyObject *type;
+    PyObject *val;
+    PyObject *tb;
+
+    if (!_PyArg_CheckPositional("__exit__", nargs, 3, 3)) {
+        goto exit;
+    }
+    type = args[0];
+    val = args[1];
+    tb = args[2];
+    return_value = token_exit_impl((PyContextToken *)self, type, val, tb);
+
+exit:
+    return return_value;
+}
+/*[clinic end generated code: output=01987cdbf68a951a input=a9049054013a1b77]*/
diff --git a/Python/context.c b/Python/context.c
index bb1aa42b9c5..dfdde7d1fa7 100644
--- a/Python/context.c
+++ b/Python/context.c
@@ -1231,9 +1231,47 @@ static PyGetSetDef PyContextTokenType_getsetlist[] = {
     {NULL}
 };
 
+/*[clinic input]
+_contextvars.Token.__enter__ as token_enter
+
+Enter into Token context manager.
+[clinic start generated code]*/
+
+static PyObject *
+token_enter_impl(PyContextToken *self)
+/*[clinic end generated code: output=9af4d2054e93fb75 input=41a3d6c4195fd47a]*/
+{
+    return Py_NewRef(self);
+}
+
+/*[clinic input]
+_contextvars.Token.__exit__ as token_exit
+
+    type: object
+    val: object
+    tb: object
+    /
+
+Exit from Token context manager, restore the linked ContextVar.
+[clinic start generated code]*/
+
+static PyObject *
+token_exit_impl(PyContextToken *self, PyObject *type, PyObject *val,
+                PyObject *tb)
+/*[clinic end generated code: output=3e6a1c95d3da703a input=7f117445f0ccd92e]*/
+{
+    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")},
+    TOKEN_ENTER_METHODDEF
+    TOKEN_EXIT_METHODDEF
     {NULL}
 };