xref: /third_party/python/Python/context.c (revision 7db96d56)
1#include "Python.h"
2#include "pycore_call.h"          // _PyObject_VectorcallTstate()
3#include "pycore_context.h"
4#include "pycore_gc.h"            // _PyObject_GC_MAY_BE_TRACKED()
5#include "pycore_hamt.h"
6#include "pycore_initconfig.h"    // _PyStatus_OK()
7#include "pycore_object.h"
8#include "pycore_pyerrors.h"
9#include "pycore_pystate.h"       // _PyThreadState_GET()
10#include "structmember.h"         // PyMemberDef
11
12
13#include "clinic/context.c.h"
14/*[clinic input]
15module _contextvars
16[clinic start generated code]*/
17/*[clinic end generated code: output=da39a3ee5e6b4b0d input=a0955718c8b8cea6]*/
18
19
20#define ENSURE_Context(o, err_ret)                                  \
21    if (!PyContext_CheckExact(o)) {                                 \
22        PyErr_SetString(PyExc_TypeError,                            \
23                        "an instance of Context was expected");     \
24        return err_ret;                                             \
25    }
26
27#define ENSURE_ContextVar(o, err_ret)                               \
28    if (!PyContextVar_CheckExact(o)) {                              \
29        PyErr_SetString(PyExc_TypeError,                            \
30                       "an instance of ContextVar was expected");   \
31        return err_ret;                                             \
32    }
33
34#define ENSURE_ContextToken(o, err_ret)                             \
35    if (!PyContextToken_CheckExact(o)) {                            \
36        PyErr_SetString(PyExc_TypeError,                            \
37                        "an instance of Token was expected");       \
38        return err_ret;                                             \
39    }
40
41
42/////////////////////////// Context API
43
44
45static PyContext *
46context_new_empty(void);
47
48static PyContext *
49context_new_from_vars(PyHamtObject *vars);
50
51static inline PyContext *
52context_get(void);
53
54static PyContextToken *
55token_new(PyContext *ctx, PyContextVar *var, PyObject *val);
56
57static PyContextVar *
58contextvar_new(PyObject *name, PyObject *def);
59
60static int
61contextvar_set(PyContextVar *var, PyObject *val);
62
63static int
64contextvar_del(PyContextVar *var);
65
66
67#if PyContext_MAXFREELIST > 0
68static struct _Py_context_state *
69get_context_state(void)
70{
71    PyInterpreterState *interp = _PyInterpreterState_GET();
72    return &interp->context;
73}
74#endif
75
76
77PyObject *
78_PyContext_NewHamtForTests(void)
79{
80    return (PyObject *)_PyHamt_New();
81}
82
83
84PyObject *
85PyContext_New(void)
86{
87    return (PyObject *)context_new_empty();
88}
89
90
91PyObject *
92PyContext_Copy(PyObject * octx)
93{
94    ENSURE_Context(octx, NULL)
95    PyContext *ctx = (PyContext *)octx;
96    return (PyObject *)context_new_from_vars(ctx->ctx_vars);
97}
98
99
100PyObject *
101PyContext_CopyCurrent(void)
102{
103    PyContext *ctx = context_get();
104    if (ctx == NULL) {
105        return NULL;
106    }
107
108    return (PyObject *)context_new_from_vars(ctx->ctx_vars);
109}
110
111
112static int
113_PyContext_Enter(PyThreadState *ts, PyObject *octx)
114{
115    ENSURE_Context(octx, -1)
116    PyContext *ctx = (PyContext *)octx;
117
118    if (ctx->ctx_entered) {
119        _PyErr_Format(ts, PyExc_RuntimeError,
120                      "cannot enter context: %R is already entered", ctx);
121        return -1;
122    }
123
124    ctx->ctx_prev = (PyContext *)ts->context;  /* borrow */
125    ctx->ctx_entered = 1;
126
127    Py_INCREF(ctx);
128    ts->context = (PyObject *)ctx;
129    ts->context_ver++;
130
131    return 0;
132}
133
134
135int
136PyContext_Enter(PyObject *octx)
137{
138    PyThreadState *ts = _PyThreadState_GET();
139    assert(ts != NULL);
140    return _PyContext_Enter(ts, octx);
141}
142
143
144static int
145_PyContext_Exit(PyThreadState *ts, PyObject *octx)
146{
147    ENSURE_Context(octx, -1)
148    PyContext *ctx = (PyContext *)octx;
149
150    if (!ctx->ctx_entered) {
151        PyErr_Format(PyExc_RuntimeError,
152                     "cannot exit context: %R has not been entered", ctx);
153        return -1;
154    }
155
156    if (ts->context != (PyObject *)ctx) {
157        /* Can only happen if someone misuses the C API */
158        PyErr_SetString(PyExc_RuntimeError,
159                        "cannot exit context: thread state references "
160                        "a different context object");
161        return -1;
162    }
163
164    Py_SETREF(ts->context, (PyObject *)ctx->ctx_prev);
165    ts->context_ver++;
166
167    ctx->ctx_prev = NULL;
168    ctx->ctx_entered = 0;
169
170    return 0;
171}
172
173int
174PyContext_Exit(PyObject *octx)
175{
176    PyThreadState *ts = _PyThreadState_GET();
177    assert(ts != NULL);
178    return _PyContext_Exit(ts, octx);
179}
180
181
182PyObject *
183PyContextVar_New(const char *name, PyObject *def)
184{
185    PyObject *pyname = PyUnicode_FromString(name);
186    if (pyname == NULL) {
187        return NULL;
188    }
189    PyContextVar *var = contextvar_new(pyname, def);
190    Py_DECREF(pyname);
191    return (PyObject *)var;
192}
193
194
195int
196PyContextVar_Get(PyObject *ovar, PyObject *def, PyObject **val)
197{
198    ENSURE_ContextVar(ovar, -1)
199    PyContextVar *var = (PyContextVar *)ovar;
200
201    PyThreadState *ts = _PyThreadState_GET();
202    assert(ts != NULL);
203    if (ts->context == NULL) {
204        goto not_found;
205    }
206
207    if (var->var_cached != NULL &&
208            var->var_cached_tsid == ts->id &&
209            var->var_cached_tsver == ts->context_ver)
210    {
211        *val = var->var_cached;
212        goto found;
213    }
214
215    assert(PyContext_CheckExact(ts->context));
216    PyHamtObject *vars = ((PyContext *)ts->context)->ctx_vars;
217
218    PyObject *found = NULL;
219    int res = _PyHamt_Find(vars, (PyObject*)var, &found);
220    if (res < 0) {
221        goto error;
222    }
223    if (res == 1) {
224        assert(found != NULL);
225        var->var_cached = found;  /* borrow */
226        var->var_cached_tsid = ts->id;
227        var->var_cached_tsver = ts->context_ver;
228
229        *val = found;
230        goto found;
231    }
232
233not_found:
234    if (def == NULL) {
235        if (var->var_default != NULL) {
236            *val = var->var_default;
237            goto found;
238        }
239
240        *val = NULL;
241        goto found;
242    }
243    else {
244        *val = def;
245        goto found;
246   }
247
248found:
249    Py_XINCREF(*val);
250    return 0;
251
252error:
253    *val = NULL;
254    return -1;
255}
256
257
258PyObject *
259PyContextVar_Set(PyObject *ovar, PyObject *val)
260{
261    ENSURE_ContextVar(ovar, NULL)
262    PyContextVar *var = (PyContextVar *)ovar;
263
264    if (!PyContextVar_CheckExact(var)) {
265        PyErr_SetString(
266            PyExc_TypeError, "an instance of ContextVar was expected");
267        return NULL;
268    }
269
270    PyContext *ctx = context_get();
271    if (ctx == NULL) {
272        return NULL;
273    }
274
275    PyObject *old_val = NULL;
276    int found = _PyHamt_Find(ctx->ctx_vars, (PyObject *)var, &old_val);
277    if (found < 0) {
278        return NULL;
279    }
280
281    Py_XINCREF(old_val);
282    PyContextToken *tok = token_new(ctx, var, old_val);
283    Py_XDECREF(old_val);
284
285    if (contextvar_set(var, val)) {
286        Py_DECREF(tok);
287        return NULL;
288    }
289
290    return (PyObject *)tok;
291}
292
293
294int
295PyContextVar_Reset(PyObject *ovar, PyObject *otok)
296{
297    ENSURE_ContextVar(ovar, -1)
298    ENSURE_ContextToken(otok, -1)
299    PyContextVar *var = (PyContextVar *)ovar;
300    PyContextToken *tok = (PyContextToken *)otok;
301
302    if (tok->tok_used) {
303        PyErr_Format(PyExc_RuntimeError,
304                     "%R has already been used once", tok);
305        return -1;
306    }
307
308    if (var != tok->tok_var) {
309        PyErr_Format(PyExc_ValueError,
310                     "%R was created by a different ContextVar", tok);
311        return -1;
312    }
313
314    PyContext *ctx = context_get();
315    if (ctx != tok->tok_ctx) {
316        PyErr_Format(PyExc_ValueError,
317                     "%R was created in a different Context", tok);
318        return -1;
319    }
320
321    tok->tok_used = 1;
322
323    if (tok->tok_oldval == NULL) {
324        return contextvar_del(var);
325    }
326    else {
327        return contextvar_set(var, tok->tok_oldval);
328    }
329}
330
331
332/////////////////////////// PyContext
333
334/*[clinic input]
335class _contextvars.Context "PyContext *" "&PyContext_Type"
336[clinic start generated code]*/
337/*[clinic end generated code: output=da39a3ee5e6b4b0d input=bdf87f8e0cb580e8]*/
338
339
340static inline PyContext *
341_context_alloc(void)
342{
343    PyContext *ctx;
344#if PyContext_MAXFREELIST > 0
345    struct _Py_context_state *state = get_context_state();
346#ifdef Py_DEBUG
347    // _context_alloc() must not be called after _PyContext_Fini()
348    assert(state->numfree != -1);
349#endif
350    if (state->numfree) {
351        state->numfree--;
352        ctx = state->freelist;
353        state->freelist = (PyContext *)ctx->ctx_weakreflist;
354        OBJECT_STAT_INC(from_freelist);
355        ctx->ctx_weakreflist = NULL;
356        _Py_NewReference((PyObject *)ctx);
357    }
358    else
359#endif
360    {
361        ctx = PyObject_GC_New(PyContext, &PyContext_Type);
362        if (ctx == NULL) {
363            return NULL;
364        }
365    }
366
367    ctx->ctx_vars = NULL;
368    ctx->ctx_prev = NULL;
369    ctx->ctx_entered = 0;
370    ctx->ctx_weakreflist = NULL;
371
372    return ctx;
373}
374
375
376static PyContext *
377context_new_empty(void)
378{
379    PyContext *ctx = _context_alloc();
380    if (ctx == NULL) {
381        return NULL;
382    }
383
384    ctx->ctx_vars = _PyHamt_New();
385    if (ctx->ctx_vars == NULL) {
386        Py_DECREF(ctx);
387        return NULL;
388    }
389
390    _PyObject_GC_TRACK(ctx);
391    return ctx;
392}
393
394
395static PyContext *
396context_new_from_vars(PyHamtObject *vars)
397{
398    PyContext *ctx = _context_alloc();
399    if (ctx == NULL) {
400        return NULL;
401    }
402
403    Py_INCREF(vars);
404    ctx->ctx_vars = vars;
405
406    _PyObject_GC_TRACK(ctx);
407    return ctx;
408}
409
410
411static inline PyContext *
412context_get(void)
413{
414    PyThreadState *ts = _PyThreadState_GET();
415    assert(ts != NULL);
416    PyContext *current_ctx = (PyContext *)ts->context;
417    if (current_ctx == NULL) {
418        current_ctx = context_new_empty();
419        if (current_ctx == NULL) {
420            return NULL;
421        }
422        ts->context = (PyObject *)current_ctx;
423    }
424    return current_ctx;
425}
426
427static int
428context_check_key_type(PyObject *key)
429{
430    if (!PyContextVar_CheckExact(key)) {
431        // abort();
432        PyErr_Format(PyExc_TypeError,
433                     "a ContextVar key was expected, got %R", key);
434        return -1;
435    }
436    return 0;
437}
438
439static PyObject *
440context_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
441{
442    if (PyTuple_Size(args) || (kwds != NULL && PyDict_Size(kwds))) {
443        PyErr_SetString(
444            PyExc_TypeError, "Context() does not accept any arguments");
445        return NULL;
446    }
447    return PyContext_New();
448}
449
450static int
451context_tp_clear(PyContext *self)
452{
453    Py_CLEAR(self->ctx_prev);
454    Py_CLEAR(self->ctx_vars);
455    return 0;
456}
457
458static int
459context_tp_traverse(PyContext *self, visitproc visit, void *arg)
460{
461    Py_VISIT(self->ctx_prev);
462    Py_VISIT(self->ctx_vars);
463    return 0;
464}
465
466static void
467context_tp_dealloc(PyContext *self)
468{
469    _PyObject_GC_UNTRACK(self);
470
471    if (self->ctx_weakreflist != NULL) {
472        PyObject_ClearWeakRefs((PyObject*)self);
473    }
474    (void)context_tp_clear(self);
475
476#if PyContext_MAXFREELIST > 0
477    struct _Py_context_state *state = get_context_state();
478#ifdef Py_DEBUG
479    // _context_alloc() must not be called after _PyContext_Fini()
480    assert(state->numfree != -1);
481#endif
482    if (state->numfree < PyContext_MAXFREELIST) {
483        state->numfree++;
484        self->ctx_weakreflist = (PyObject *)state->freelist;
485        state->freelist = self;
486        OBJECT_STAT_INC(to_freelist);
487    }
488    else
489#endif
490    {
491        Py_TYPE(self)->tp_free(self);
492    }
493}
494
495static PyObject *
496context_tp_iter(PyContext *self)
497{
498    return _PyHamt_NewIterKeys(self->ctx_vars);
499}
500
501static PyObject *
502context_tp_richcompare(PyObject *v, PyObject *w, int op)
503{
504    if (!PyContext_CheckExact(v) || !PyContext_CheckExact(w) ||
505            (op != Py_EQ && op != Py_NE))
506    {
507        Py_RETURN_NOTIMPLEMENTED;
508    }
509
510    int res = _PyHamt_Eq(
511        ((PyContext *)v)->ctx_vars, ((PyContext *)w)->ctx_vars);
512    if (res < 0) {
513        return NULL;
514    }
515
516    if (op == Py_NE) {
517        res = !res;
518    }
519
520    if (res) {
521        Py_RETURN_TRUE;
522    }
523    else {
524        Py_RETURN_FALSE;
525    }
526}
527
528static Py_ssize_t
529context_tp_len(PyContext *self)
530{
531    return _PyHamt_Len(self->ctx_vars);
532}
533
534static PyObject *
535context_tp_subscript(PyContext *self, PyObject *key)
536{
537    if (context_check_key_type(key)) {
538        return NULL;
539    }
540    PyObject *val = NULL;
541    int found = _PyHamt_Find(self->ctx_vars, key, &val);
542    if (found < 0) {
543        return NULL;
544    }
545    if (found == 0) {
546        PyErr_SetObject(PyExc_KeyError, key);
547        return NULL;
548    }
549    Py_INCREF(val);
550    return val;
551}
552
553static int
554context_tp_contains(PyContext *self, PyObject *key)
555{
556    if (context_check_key_type(key)) {
557        return -1;
558    }
559    PyObject *val = NULL;
560    return _PyHamt_Find(self->ctx_vars, key, &val);
561}
562
563
564/*[clinic input]
565_contextvars.Context.get
566    key: object
567    default: object = None
568    /
569
570Return the value for `key` if `key` has the value in the context object.
571
572If `key` does not exist, return `default`. If `default` is not given,
573return None.
574[clinic start generated code]*/
575
576static PyObject *
577_contextvars_Context_get_impl(PyContext *self, PyObject *key,
578                              PyObject *default_value)
579/*[clinic end generated code: output=0c54aa7664268189 input=c8eeb81505023995]*/
580{
581    if (context_check_key_type(key)) {
582        return NULL;
583    }
584
585    PyObject *val = NULL;
586    int found = _PyHamt_Find(self->ctx_vars, key, &val);
587    if (found < 0) {
588        return NULL;
589    }
590    if (found == 0) {
591        Py_INCREF(default_value);
592        return default_value;
593    }
594    Py_INCREF(val);
595    return val;
596}
597
598
599/*[clinic input]
600_contextvars.Context.items
601
602Return all variables and their values in the context object.
603
604The result is returned as a list of 2-tuples (variable, value).
605[clinic start generated code]*/
606
607static PyObject *
608_contextvars_Context_items_impl(PyContext *self)
609/*[clinic end generated code: output=fa1655c8a08502af input=00db64ae379f9f42]*/
610{
611    return _PyHamt_NewIterItems(self->ctx_vars);
612}
613
614
615/*[clinic input]
616_contextvars.Context.keys
617
618Return a list of all variables in the context object.
619[clinic start generated code]*/
620
621static PyObject *
622_contextvars_Context_keys_impl(PyContext *self)
623/*[clinic end generated code: output=177227c6b63ec0e2 input=114b53aebca3449c]*/
624{
625    return _PyHamt_NewIterKeys(self->ctx_vars);
626}
627
628
629/*[clinic input]
630_contextvars.Context.values
631
632Return a list of all variables' values in the context object.
633[clinic start generated code]*/
634
635static PyObject *
636_contextvars_Context_values_impl(PyContext *self)
637/*[clinic end generated code: output=d286dabfc8db6dde input=ce8075d04a6ea526]*/
638{
639    return _PyHamt_NewIterValues(self->ctx_vars);
640}
641
642
643/*[clinic input]
644_contextvars.Context.copy
645
646Return a shallow copy of the context object.
647[clinic start generated code]*/
648
649static PyObject *
650_contextvars_Context_copy_impl(PyContext *self)
651/*[clinic end generated code: output=30ba8896c4707a15 input=ebafdbdd9c72d592]*/
652{
653    return (PyObject *)context_new_from_vars(self->ctx_vars);
654}
655
656
657static PyObject *
658context_run(PyContext *self, PyObject *const *args,
659            Py_ssize_t nargs, PyObject *kwnames)
660{
661    PyThreadState *ts = _PyThreadState_GET();
662
663    if (nargs < 1) {
664        _PyErr_SetString(ts, PyExc_TypeError,
665                         "run() missing 1 required positional argument");
666        return NULL;
667    }
668
669    if (_PyContext_Enter(ts, (PyObject *)self)) {
670        return NULL;
671    }
672
673    PyObject *call_result = _PyObject_VectorcallTstate(
674        ts, args[0], args + 1, nargs - 1, kwnames);
675
676    if (_PyContext_Exit(ts, (PyObject *)self)) {
677        return NULL;
678    }
679
680    return call_result;
681}
682
683
684static PyMethodDef PyContext_methods[] = {
685    _CONTEXTVARS_CONTEXT_GET_METHODDEF
686    _CONTEXTVARS_CONTEXT_ITEMS_METHODDEF
687    _CONTEXTVARS_CONTEXT_KEYS_METHODDEF
688    _CONTEXTVARS_CONTEXT_VALUES_METHODDEF
689    _CONTEXTVARS_CONTEXT_COPY_METHODDEF
690    {"run", _PyCFunction_CAST(context_run), METH_FASTCALL | METH_KEYWORDS, NULL},
691    {NULL, NULL}
692};
693
694static PySequenceMethods PyContext_as_sequence = {
695    0,                                   /* sq_length */
696    0,                                   /* sq_concat */
697    0,                                   /* sq_repeat */
698    0,                                   /* sq_item */
699    0,                                   /* sq_slice */
700    0,                                   /* sq_ass_item */
701    0,                                   /* sq_ass_slice */
702    (objobjproc)context_tp_contains,     /* sq_contains */
703    0,                                   /* sq_inplace_concat */
704    0,                                   /* sq_inplace_repeat */
705};
706
707static PyMappingMethods PyContext_as_mapping = {
708    (lenfunc)context_tp_len,             /* mp_length */
709    (binaryfunc)context_tp_subscript,    /* mp_subscript */
710};
711
712PyTypeObject PyContext_Type = {
713    PyVarObject_HEAD_INIT(&PyType_Type, 0)
714    "_contextvars.Context",
715    sizeof(PyContext),
716    .tp_methods = PyContext_methods,
717    .tp_as_mapping = &PyContext_as_mapping,
718    .tp_as_sequence = &PyContext_as_sequence,
719    .tp_iter = (getiterfunc)context_tp_iter,
720    .tp_dealloc = (destructor)context_tp_dealloc,
721    .tp_getattro = PyObject_GenericGetAttr,
722    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
723    .tp_richcompare = context_tp_richcompare,
724    .tp_traverse = (traverseproc)context_tp_traverse,
725    .tp_clear = (inquiry)context_tp_clear,
726    .tp_new = context_tp_new,
727    .tp_weaklistoffset = offsetof(PyContext, ctx_weakreflist),
728    .tp_hash = PyObject_HashNotImplemented,
729};
730
731
732/////////////////////////// ContextVar
733
734
735static int
736contextvar_set(PyContextVar *var, PyObject *val)
737{
738    var->var_cached = NULL;
739    PyThreadState *ts = _PyThreadState_GET();
740
741    PyContext *ctx = context_get();
742    if (ctx == NULL) {
743        return -1;
744    }
745
746    PyHamtObject *new_vars = _PyHamt_Assoc(
747        ctx->ctx_vars, (PyObject *)var, val);
748    if (new_vars == NULL) {
749        return -1;
750    }
751
752    Py_SETREF(ctx->ctx_vars, new_vars);
753
754    var->var_cached = val;  /* borrow */
755    var->var_cached_tsid = ts->id;
756    var->var_cached_tsver = ts->context_ver;
757    return 0;
758}
759
760static int
761contextvar_del(PyContextVar *var)
762{
763    var->var_cached = NULL;
764
765    PyContext *ctx = context_get();
766    if (ctx == NULL) {
767        return -1;
768    }
769
770    PyHamtObject *vars = ctx->ctx_vars;
771    PyHamtObject *new_vars = _PyHamt_Without(vars, (PyObject *)var);
772    if (new_vars == NULL) {
773        return -1;
774    }
775
776    if (vars == new_vars) {
777        Py_DECREF(new_vars);
778        PyErr_SetObject(PyExc_LookupError, (PyObject *)var);
779        return -1;
780    }
781
782    Py_SETREF(ctx->ctx_vars, new_vars);
783    return 0;
784}
785
786static Py_hash_t
787contextvar_generate_hash(void *addr, PyObject *name)
788{
789    /* Take hash of `name` and XOR it with the object's addr.
790
791       The structure of the tree is encoded in objects' hashes, which
792       means that sufficiently similar hashes would result in tall trees
793       with many Collision nodes.  Which would, in turn, result in slower
794       get and set operations.
795
796       The XORing helps to ensure that:
797
798       (1) sequentially allocated ContextVar objects have
799           different hashes;
800
801       (2) context variables with equal names have
802           different hashes.
803    */
804
805    Py_hash_t name_hash = PyObject_Hash(name);
806    if (name_hash == -1) {
807        return -1;
808    }
809
810    Py_hash_t res = _Py_HashPointer(addr) ^ name_hash;
811    return res == -1 ? -2 : res;
812}
813
814static PyContextVar *
815contextvar_new(PyObject *name, PyObject *def)
816{
817    if (!PyUnicode_Check(name)) {
818        PyErr_SetString(PyExc_TypeError,
819                        "context variable name must be a str");
820        return NULL;
821    }
822
823    PyContextVar *var = PyObject_GC_New(PyContextVar, &PyContextVar_Type);
824    if (var == NULL) {
825        return NULL;
826    }
827
828    var->var_hash = contextvar_generate_hash(var, name);
829    if (var->var_hash == -1) {
830        Py_DECREF(var);
831        return NULL;
832    }
833
834    Py_INCREF(name);
835    var->var_name = name;
836
837    Py_XINCREF(def);
838    var->var_default = def;
839
840    var->var_cached = NULL;
841    var->var_cached_tsid = 0;
842    var->var_cached_tsver = 0;
843
844    if (_PyObject_GC_MAY_BE_TRACKED(name) ||
845            (def != NULL && _PyObject_GC_MAY_BE_TRACKED(def)))
846    {
847        PyObject_GC_Track(var);
848    }
849    return var;
850}
851
852
853/*[clinic input]
854class _contextvars.ContextVar "PyContextVar *" "&PyContextVar_Type"
855[clinic start generated code]*/
856/*[clinic end generated code: output=da39a3ee5e6b4b0d input=445da935fa8883c3]*/
857
858
859static PyObject *
860contextvar_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
861{
862    static char *kwlist[] = {"", "default", NULL};
863    PyObject *name;
864    PyObject *def = NULL;
865
866    if (!PyArg_ParseTupleAndKeywords(
867            args, kwds, "O|$O:ContextVar", kwlist, &name, &def))
868    {
869        return NULL;
870    }
871
872    return (PyObject *)contextvar_new(name, def);
873}
874
875static int
876contextvar_tp_clear(PyContextVar *self)
877{
878    Py_CLEAR(self->var_name);
879    Py_CLEAR(self->var_default);
880    self->var_cached = NULL;
881    self->var_cached_tsid = 0;
882    self->var_cached_tsver = 0;
883    return 0;
884}
885
886static int
887contextvar_tp_traverse(PyContextVar *self, visitproc visit, void *arg)
888{
889    Py_VISIT(self->var_name);
890    Py_VISIT(self->var_default);
891    return 0;
892}
893
894static void
895contextvar_tp_dealloc(PyContextVar *self)
896{
897    PyObject_GC_UnTrack(self);
898    (void)contextvar_tp_clear(self);
899    Py_TYPE(self)->tp_free(self);
900}
901
902static Py_hash_t
903contextvar_tp_hash(PyContextVar *self)
904{
905    return self->var_hash;
906}
907
908static PyObject *
909contextvar_tp_repr(PyContextVar *self)
910{
911    _PyUnicodeWriter writer;
912
913    _PyUnicodeWriter_Init(&writer);
914
915    if (_PyUnicodeWriter_WriteASCIIString(
916            &writer, "<ContextVar name=", 17) < 0)
917    {
918        goto error;
919    }
920
921    PyObject *name = PyObject_Repr(self->var_name);
922    if (name == NULL) {
923        goto error;
924    }
925    if (_PyUnicodeWriter_WriteStr(&writer, name) < 0) {
926        Py_DECREF(name);
927        goto error;
928    }
929    Py_DECREF(name);
930
931    if (self->var_default != NULL) {
932        if (_PyUnicodeWriter_WriteASCIIString(&writer, " default=", 9) < 0) {
933            goto error;
934        }
935
936        PyObject *def = PyObject_Repr(self->var_default);
937        if (def == NULL) {
938            goto error;
939        }
940        if (_PyUnicodeWriter_WriteStr(&writer, def) < 0) {
941            Py_DECREF(def);
942            goto error;
943        }
944        Py_DECREF(def);
945    }
946
947    PyObject *addr = PyUnicode_FromFormat(" at %p>", self);
948    if (addr == NULL) {
949        goto error;
950    }
951    if (_PyUnicodeWriter_WriteStr(&writer, addr) < 0) {
952        Py_DECREF(addr);
953        goto error;
954    }
955    Py_DECREF(addr);
956
957    return _PyUnicodeWriter_Finish(&writer);
958
959error:
960    _PyUnicodeWriter_Dealloc(&writer);
961    return NULL;
962}
963
964
965/*[clinic input]
966_contextvars.ContextVar.get
967    default: object = NULL
968    /
969
970Return a value for the context variable for the current context.
971
972If there is no value for the variable in the current context, the method will:
973 * return the value of the default argument of the method, if provided; or
974 * return the default value for the context variable, if it was created
975   with one; or
976 * raise a LookupError.
977[clinic start generated code]*/
978
979static PyObject *
980_contextvars_ContextVar_get_impl(PyContextVar *self, PyObject *default_value)
981/*[clinic end generated code: output=0746bd0aa2ced7bf input=30aa2ab9e433e401]*/
982{
983    if (!PyContextVar_CheckExact(self)) {
984        PyErr_SetString(
985            PyExc_TypeError, "an instance of ContextVar was expected");
986        return NULL;
987    }
988
989    PyObject *val;
990    if (PyContextVar_Get((PyObject *)self, default_value, &val) < 0) {
991        return NULL;
992    }
993
994    if (val == NULL) {
995        PyErr_SetObject(PyExc_LookupError, (PyObject *)self);
996        return NULL;
997    }
998
999    return val;
1000}
1001
1002/*[clinic input]
1003_contextvars.ContextVar.set
1004    value: object
1005    /
1006
1007Call to set a new value for the context variable in the current context.
1008
1009The required value argument is the new value for the context variable.
1010
1011Returns a Token object that can be used to restore the variable to its previous
1012value via the `ContextVar.reset()` method.
1013[clinic start generated code]*/
1014
1015static PyObject *
1016_contextvars_ContextVar_set(PyContextVar *self, PyObject *value)
1017/*[clinic end generated code: output=446ed5e820d6d60b input=c0a6887154227453]*/
1018{
1019    return PyContextVar_Set((PyObject *)self, value);
1020}
1021
1022/*[clinic input]
1023_contextvars.ContextVar.reset
1024    token: object
1025    /
1026
1027Reset the context variable.
1028
1029The variable is reset to the value it had before the `ContextVar.set()` that
1030created the token was used.
1031[clinic start generated code]*/
1032
1033static PyObject *
1034_contextvars_ContextVar_reset(PyContextVar *self, PyObject *token)
1035/*[clinic end generated code: output=d4ee34d0742d62ee input=ebe2881e5af4ffda]*/
1036{
1037    if (!PyContextToken_CheckExact(token)) {
1038        PyErr_Format(PyExc_TypeError,
1039                     "expected an instance of Token, got %R", token);
1040        return NULL;
1041    }
1042
1043    if (PyContextVar_Reset((PyObject *)self, token)) {
1044        return NULL;
1045    }
1046
1047    Py_RETURN_NONE;
1048}
1049
1050
1051static PyMemberDef PyContextVar_members[] = {
1052    {"name", T_OBJECT, offsetof(PyContextVar, var_name), READONLY},
1053    {NULL}
1054};
1055
1056static PyMethodDef PyContextVar_methods[] = {
1057    _CONTEXTVARS_CONTEXTVAR_GET_METHODDEF
1058    _CONTEXTVARS_CONTEXTVAR_SET_METHODDEF
1059    _CONTEXTVARS_CONTEXTVAR_RESET_METHODDEF
1060    {"__class_getitem__", Py_GenericAlias,
1061    METH_O|METH_CLASS,       PyDoc_STR("See PEP 585")},
1062    {NULL, NULL}
1063};
1064
1065PyTypeObject PyContextVar_Type = {
1066    PyVarObject_HEAD_INIT(&PyType_Type, 0)
1067    "_contextvars.ContextVar",
1068    sizeof(PyContextVar),
1069    .tp_methods = PyContextVar_methods,
1070    .tp_members = PyContextVar_members,
1071    .tp_dealloc = (destructor)contextvar_tp_dealloc,
1072    .tp_getattro = PyObject_GenericGetAttr,
1073    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
1074    .tp_traverse = (traverseproc)contextvar_tp_traverse,
1075    .tp_clear = (inquiry)contextvar_tp_clear,
1076    .tp_new = contextvar_tp_new,
1077    .tp_free = PyObject_GC_Del,
1078    .tp_hash = (hashfunc)contextvar_tp_hash,
1079    .tp_repr = (reprfunc)contextvar_tp_repr,
1080};
1081
1082
1083/////////////////////////// Token
1084
1085static PyObject * get_token_missing(void);
1086
1087
1088/*[clinic input]
1089class _contextvars.Token "PyContextToken *" "&PyContextToken_Type"
1090[clinic start generated code]*/
1091/*[clinic end generated code: output=da39a3ee5e6b4b0d input=338a5e2db13d3f5b]*/
1092
1093
1094static PyObject *
1095token_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1096{
1097    PyErr_SetString(PyExc_RuntimeError,
1098                    "Tokens can only be created by ContextVars");
1099    return NULL;
1100}
1101
1102static int
1103token_tp_clear(PyContextToken *self)
1104{
1105    Py_CLEAR(self->tok_ctx);
1106    Py_CLEAR(self->tok_var);
1107    Py_CLEAR(self->tok_oldval);
1108    return 0;
1109}
1110
1111static int
1112token_tp_traverse(PyContextToken *self, visitproc visit, void *arg)
1113{
1114    Py_VISIT(self->tok_ctx);
1115    Py_VISIT(self->tok_var);
1116    Py_VISIT(self->tok_oldval);
1117    return 0;
1118}
1119
1120static void
1121token_tp_dealloc(PyContextToken *self)
1122{
1123    PyObject_GC_UnTrack(self);
1124    (void)token_tp_clear(self);
1125    Py_TYPE(self)->tp_free(self);
1126}
1127
1128static PyObject *
1129token_tp_repr(PyContextToken *self)
1130{
1131    _PyUnicodeWriter writer;
1132
1133    _PyUnicodeWriter_Init(&writer);
1134
1135    if (_PyUnicodeWriter_WriteASCIIString(&writer, "<Token", 6) < 0) {
1136        goto error;
1137    }
1138
1139    if (self->tok_used) {
1140        if (_PyUnicodeWriter_WriteASCIIString(&writer, " used", 5) < 0) {
1141            goto error;
1142        }
1143    }
1144
1145    if (_PyUnicodeWriter_WriteASCIIString(&writer, " var=", 5) < 0) {
1146        goto error;
1147    }
1148
1149    PyObject *var = PyObject_Repr((PyObject *)self->tok_var);
1150    if (var == NULL) {
1151        goto error;
1152    }
1153    if (_PyUnicodeWriter_WriteStr(&writer, var) < 0) {
1154        Py_DECREF(var);
1155        goto error;
1156    }
1157    Py_DECREF(var);
1158
1159    PyObject *addr = PyUnicode_FromFormat(" at %p>", self);
1160    if (addr == NULL) {
1161        goto error;
1162    }
1163    if (_PyUnicodeWriter_WriteStr(&writer, addr) < 0) {
1164        Py_DECREF(addr);
1165        goto error;
1166    }
1167    Py_DECREF(addr);
1168
1169    return _PyUnicodeWriter_Finish(&writer);
1170
1171error:
1172    _PyUnicodeWriter_Dealloc(&writer);
1173    return NULL;
1174}
1175
1176static PyObject *
1177token_get_var(PyContextToken *self, void *Py_UNUSED(ignored))
1178{
1179    Py_INCREF(self->tok_var);
1180    return (PyObject *)self->tok_var;
1181}
1182
1183static PyObject *
1184token_get_old_value(PyContextToken *self, void *Py_UNUSED(ignored))
1185{
1186    if (self->tok_oldval == NULL) {
1187        return get_token_missing();
1188    }
1189
1190    Py_INCREF(self->tok_oldval);
1191    return self->tok_oldval;
1192}
1193
1194static PyGetSetDef PyContextTokenType_getsetlist[] = {
1195    {"var", (getter)token_get_var, NULL, NULL},
1196    {"old_value", (getter)token_get_old_value, NULL, NULL},
1197    {NULL}
1198};
1199
1200static PyMethodDef PyContextTokenType_methods[] = {
1201    {"__class_getitem__",    Py_GenericAlias,
1202    METH_O|METH_CLASS,       PyDoc_STR("See PEP 585")},
1203    {NULL}
1204};
1205
1206PyTypeObject PyContextToken_Type = {
1207    PyVarObject_HEAD_INIT(&PyType_Type, 0)
1208    "_contextvars.Token",
1209    sizeof(PyContextToken),
1210    .tp_methods = PyContextTokenType_methods,
1211    .tp_getset = PyContextTokenType_getsetlist,
1212    .tp_dealloc = (destructor)token_tp_dealloc,
1213    .tp_getattro = PyObject_GenericGetAttr,
1214    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
1215    .tp_traverse = (traverseproc)token_tp_traverse,
1216    .tp_clear = (inquiry)token_tp_clear,
1217    .tp_new = token_tp_new,
1218    .tp_free = PyObject_GC_Del,
1219    .tp_hash = PyObject_HashNotImplemented,
1220    .tp_repr = (reprfunc)token_tp_repr,
1221};
1222
1223static PyContextToken *
1224token_new(PyContext *ctx, PyContextVar *var, PyObject *val)
1225{
1226    PyContextToken *tok = PyObject_GC_New(PyContextToken, &PyContextToken_Type);
1227    if (tok == NULL) {
1228        return NULL;
1229    }
1230
1231    Py_INCREF(ctx);
1232    tok->tok_ctx = ctx;
1233
1234    Py_INCREF(var);
1235    tok->tok_var = var;
1236
1237    Py_XINCREF(val);
1238    tok->tok_oldval = val;
1239
1240    tok->tok_used = 0;
1241
1242    PyObject_GC_Track(tok);
1243    return tok;
1244}
1245
1246
1247/////////////////////////// Token.MISSING
1248
1249
1250static PyObject *_token_missing;
1251
1252
1253typedef struct {
1254    PyObject_HEAD
1255} PyContextTokenMissing;
1256
1257
1258static PyObject *
1259context_token_missing_tp_repr(PyObject *self)
1260{
1261    return PyUnicode_FromString("<Token.MISSING>");
1262}
1263
1264
1265PyTypeObject _PyContextTokenMissing_Type = {
1266    PyVarObject_HEAD_INIT(&PyType_Type, 0)
1267    "Token.MISSING",
1268    sizeof(PyContextTokenMissing),
1269    .tp_getattro = PyObject_GenericGetAttr,
1270    .tp_flags = Py_TPFLAGS_DEFAULT,
1271    .tp_repr = context_token_missing_tp_repr,
1272};
1273
1274
1275static PyObject *
1276get_token_missing(void)
1277{
1278    if (_token_missing != NULL) {
1279        Py_INCREF(_token_missing);
1280        return _token_missing;
1281    }
1282
1283    _token_missing = (PyObject *)PyObject_New(
1284        PyContextTokenMissing, &_PyContextTokenMissing_Type);
1285    if (_token_missing == NULL) {
1286        return NULL;
1287    }
1288
1289    Py_INCREF(_token_missing);
1290    return _token_missing;
1291}
1292
1293
1294///////////////////////////
1295
1296
1297void
1298_PyContext_ClearFreeList(PyInterpreterState *interp)
1299{
1300#if PyContext_MAXFREELIST > 0
1301    struct _Py_context_state *state = &interp->context;
1302    for (; state->numfree; state->numfree--) {
1303        PyContext *ctx = state->freelist;
1304        state->freelist = (PyContext *)ctx->ctx_weakreflist;
1305        ctx->ctx_weakreflist = NULL;
1306        PyObject_GC_Del(ctx);
1307    }
1308#endif
1309}
1310
1311
1312void
1313_PyContext_Fini(PyInterpreterState *interp)
1314{
1315    if (_Py_IsMainInterpreter(interp)) {
1316        Py_CLEAR(_token_missing);
1317    }
1318    _PyContext_ClearFreeList(interp);
1319#if defined(Py_DEBUG) && PyContext_MAXFREELIST > 0
1320    struct _Py_context_state *state = &interp->context;
1321    state->numfree = -1;
1322#endif
1323    _PyHamt_Fini(interp);
1324}
1325
1326
1327PyStatus
1328_PyContext_Init(PyInterpreterState *interp)
1329{
1330    if (!_Py_IsMainInterpreter(interp)) {
1331        return _PyStatus_OK();
1332    }
1333
1334    PyObject *missing = get_token_missing();
1335    if (PyDict_SetItemString(
1336        PyContextToken_Type.tp_dict, "MISSING", missing))
1337    {
1338        Py_DECREF(missing);
1339        return _PyStatus_ERR("can't init context types");
1340    }
1341    Py_DECREF(missing);
1342
1343    return _PyStatus_OK();
1344}
1345