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
11
12
13 #include "clinic/context.c.h"
14 /*[clinic input]
15 module _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
45 static PyContext *
46 context_new_empty(void);
47
48 static PyContext *
49 context_new_from_vars(PyHamtObject *vars);
50
51 static inline PyContext *
52 context_get(void);
53
54 static PyContextToken *
55 token_new(PyContext *ctx, PyContextVar *var, PyObject *val);
56
57 static PyContextVar *
58 contextvar_new(PyObject *name, PyObject *def);
59
60 static int
61 contextvar_set(PyContextVar *var, PyObject *val);
62
63 static int
64 contextvar_del(PyContextVar *var);
65
66
67 #ifdef WITH_FREELISTS
68 static struct _Py_context_freelist *
get_context_freelist(void)69 get_context_freelist(void)
70 {
71 struct _Py_object_freelists *freelists = _Py_object_freelists_GET();
72 return &freelists->contexts;
73 }
74 #endif
75
76
77 PyObject *
_PyContext_NewHamtForTests(void)78 _PyContext_NewHamtForTests(void)
79 {
80 return (PyObject *)_PyHamt_New();
81 }
82
83
84 PyObject *
PyContext_New(void)85 PyContext_New(void)
86 {
87 return (PyObject *)context_new_empty();
88 }
89
90
91 PyObject *
PyContext_Copy(PyObject * octx)92 PyContext_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
100 PyObject *
PyContext_CopyCurrent(void)101 PyContext_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
112 static int
_PyContext_Enter(PyThreadState * ts,PyObject * octx)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 ts->context = Py_NewRef(ctx);
128 ts->context_ver++;
129
130 return 0;
131 }
132
133
134 int
PyContext_Enter(PyObject * octx)135 PyContext_Enter(PyObject *octx)
136 {
137 PyThreadState *ts = _PyThreadState_GET();
138 assert(ts != NULL);
139 return _PyContext_Enter(ts, octx);
140 }
141
142
143 static int
_PyContext_Exit(PyThreadState * ts,PyObject * octx)144 _PyContext_Exit(PyThreadState *ts, PyObject *octx)
145 {
146 ENSURE_Context(octx, -1)
147 PyContext *ctx = (PyContext *)octx;
148
149 if (!ctx->ctx_entered) {
150 PyErr_Format(PyExc_RuntimeError,
151 "cannot exit context: %R has not been entered", ctx);
152 return -1;
153 }
154
155 if (ts->context != (PyObject *)ctx) {
156 /* Can only happen if someone misuses the C API */
157 PyErr_SetString(PyExc_RuntimeError,
158 "cannot exit context: thread state references "
159 "a different context object");
160 return -1;
161 }
162
163 Py_SETREF(ts->context, (PyObject *)ctx->ctx_prev);
164 ts->context_ver++;
165
166 ctx->ctx_prev = NULL;
167 ctx->ctx_entered = 0;
168
169 return 0;
170 }
171
172 int
PyContext_Exit(PyObject * octx)173 PyContext_Exit(PyObject *octx)
174 {
175 PyThreadState *ts = _PyThreadState_GET();
176 assert(ts != NULL);
177 return _PyContext_Exit(ts, octx);
178 }
179
180
181 PyObject *
PyContextVar_New(const char * name,PyObject * def)182 PyContextVar_New(const char *name, PyObject *def)
183 {
184 PyObject *pyname = PyUnicode_FromString(name);
185 if (pyname == NULL) {
186 return NULL;
187 }
188 PyContextVar *var = contextvar_new(pyname, def);
189 Py_DECREF(pyname);
190 return (PyObject *)var;
191 }
192
193
194 int
PyContextVar_Get(PyObject * ovar,PyObject * def,PyObject ** val)195 PyContextVar_Get(PyObject *ovar, PyObject *def, PyObject **val)
196 {
197 ENSURE_ContextVar(ovar, -1)
198 PyContextVar *var = (PyContextVar *)ovar;
199
200 PyThreadState *ts = _PyThreadState_GET();
201 assert(ts != NULL);
202 if (ts->context == NULL) {
203 goto not_found;
204 }
205
206 #ifndef Py_GIL_DISABLED
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 #endif
215
216 assert(PyContext_CheckExact(ts->context));
217 PyHamtObject *vars = ((PyContext *)ts->context)->ctx_vars;
218
219 PyObject *found = NULL;
220 int res = _PyHamt_Find(vars, (PyObject*)var, &found);
221 if (res < 0) {
222 goto error;
223 }
224 if (res == 1) {
225 assert(found != NULL);
226 #ifndef Py_GIL_DISABLED
227 var->var_cached = found; /* borrow */
228 var->var_cached_tsid = ts->id;
229 var->var_cached_tsver = ts->context_ver;
230 #endif
231
232 *val = found;
233 goto found;
234 }
235
236 not_found:
237 if (def == NULL) {
238 if (var->var_default != NULL) {
239 *val = var->var_default;
240 goto found;
241 }
242
243 *val = NULL;
244 goto found;
245 }
246 else {
247 *val = def;
248 goto found;
249 }
250
251 found:
252 Py_XINCREF(*val);
253 return 0;
254
255 error:
256 *val = NULL;
257 return -1;
258 }
259
260
261 PyObject *
PyContextVar_Set(PyObject * ovar,PyObject * val)262 PyContextVar_Set(PyObject *ovar, PyObject *val)
263 {
264 ENSURE_ContextVar(ovar, NULL)
265 PyContextVar *var = (PyContextVar *)ovar;
266
267 if (!PyContextVar_CheckExact(var)) {
268 PyErr_SetString(
269 PyExc_TypeError, "an instance of ContextVar was expected");
270 return NULL;
271 }
272
273 PyContext *ctx = context_get();
274 if (ctx == NULL) {
275 return NULL;
276 }
277
278 PyObject *old_val = NULL;
279 int found = _PyHamt_Find(ctx->ctx_vars, (PyObject *)var, &old_val);
280 if (found < 0) {
281 return NULL;
282 }
283
284 Py_XINCREF(old_val);
285 PyContextToken *tok = token_new(ctx, var, old_val);
286 Py_XDECREF(old_val);
287
288 if (contextvar_set(var, val)) {
289 Py_DECREF(tok);
290 return NULL;
291 }
292
293 return (PyObject *)tok;
294 }
295
296
297 int
PyContextVar_Reset(PyObject * ovar,PyObject * otok)298 PyContextVar_Reset(PyObject *ovar, PyObject *otok)
299 {
300 ENSURE_ContextVar(ovar, -1)
301 ENSURE_ContextToken(otok, -1)
302 PyContextVar *var = (PyContextVar *)ovar;
303 PyContextToken *tok = (PyContextToken *)otok;
304
305 if (tok->tok_used) {
306 PyErr_Format(PyExc_RuntimeError,
307 "%R has already been used once", tok);
308 return -1;
309 }
310
311 if (var != tok->tok_var) {
312 PyErr_Format(PyExc_ValueError,
313 "%R was created by a different ContextVar", tok);
314 return -1;
315 }
316
317 PyContext *ctx = context_get();
318 if (ctx != tok->tok_ctx) {
319 PyErr_Format(PyExc_ValueError,
320 "%R was created in a different Context", tok);
321 return -1;
322 }
323
324 tok->tok_used = 1;
325
326 if (tok->tok_oldval == NULL) {
327 return contextvar_del(var);
328 }
329 else {
330 return contextvar_set(var, tok->tok_oldval);
331 }
332 }
333
334
335 /////////////////////////// PyContext
336
337 /*[clinic input]
338 class _contextvars.Context "PyContext *" "&PyContext_Type"
339 [clinic start generated code]*/
340 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=bdf87f8e0cb580e8]*/
341
342
343 static inline PyContext *
_context_alloc(void)344 _context_alloc(void)
345 {
346 PyContext *ctx;
347 #ifdef WITH_FREELISTS
348 struct _Py_context_freelist *context_freelist = get_context_freelist();
349 if (context_freelist->numfree > 0) {
350 context_freelist->numfree--;
351 ctx = context_freelist->items;
352 context_freelist->items = (PyContext *)ctx->ctx_weakreflist;
353 OBJECT_STAT_INC(from_freelist);
354 ctx->ctx_weakreflist = NULL;
355 _Py_NewReference((PyObject *)ctx);
356 }
357 else
358 #endif
359 {
360 ctx = PyObject_GC_New(PyContext, &PyContext_Type);
361 if (ctx == NULL) {
362 return NULL;
363 }
364 }
365
366 ctx->ctx_vars = NULL;
367 ctx->ctx_prev = NULL;
368 ctx->ctx_entered = 0;
369 ctx->ctx_weakreflist = NULL;
370
371 return ctx;
372 }
373
374
375 static PyContext *
context_new_empty(void)376 context_new_empty(void)
377 {
378 PyContext *ctx = _context_alloc();
379 if (ctx == NULL) {
380 return NULL;
381 }
382
383 ctx->ctx_vars = _PyHamt_New();
384 if (ctx->ctx_vars == NULL) {
385 Py_DECREF(ctx);
386 return NULL;
387 }
388
389 _PyObject_GC_TRACK(ctx);
390 return ctx;
391 }
392
393
394 static PyContext *
context_new_from_vars(PyHamtObject * vars)395 context_new_from_vars(PyHamtObject *vars)
396 {
397 PyContext *ctx = _context_alloc();
398 if (ctx == NULL) {
399 return NULL;
400 }
401
402 ctx->ctx_vars = (PyHamtObject*)Py_NewRef(vars);
403
404 _PyObject_GC_TRACK(ctx);
405 return ctx;
406 }
407
408
409 static inline PyContext *
context_get(void)410 context_get(void)
411 {
412 PyThreadState *ts = _PyThreadState_GET();
413 assert(ts != NULL);
414 PyContext *current_ctx = (PyContext *)ts->context;
415 if (current_ctx == NULL) {
416 current_ctx = context_new_empty();
417 if (current_ctx == NULL) {
418 return NULL;
419 }
420 ts->context = (PyObject *)current_ctx;
421 }
422 return current_ctx;
423 }
424
425 static int
context_check_key_type(PyObject * key)426 context_check_key_type(PyObject *key)
427 {
428 if (!PyContextVar_CheckExact(key)) {
429 // abort();
430 PyErr_Format(PyExc_TypeError,
431 "a ContextVar key was expected, got %R", key);
432 return -1;
433 }
434 return 0;
435 }
436
437 static PyObject *
context_tp_new(PyTypeObject * type,PyObject * args,PyObject * kwds)438 context_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
439 {
440 if (PyTuple_Size(args) || (kwds != NULL && PyDict_Size(kwds))) {
441 PyErr_SetString(
442 PyExc_TypeError, "Context() does not accept any arguments");
443 return NULL;
444 }
445 return PyContext_New();
446 }
447
448 static int
context_tp_clear(PyContext * self)449 context_tp_clear(PyContext *self)
450 {
451 Py_CLEAR(self->ctx_prev);
452 Py_CLEAR(self->ctx_vars);
453 return 0;
454 }
455
456 static int
context_tp_traverse(PyContext * self,visitproc visit,void * arg)457 context_tp_traverse(PyContext *self, visitproc visit, void *arg)
458 {
459 Py_VISIT(self->ctx_prev);
460 Py_VISIT(self->ctx_vars);
461 return 0;
462 }
463
464 static void
context_tp_dealloc(PyContext * self)465 context_tp_dealloc(PyContext *self)
466 {
467 _PyObject_GC_UNTRACK(self);
468
469 if (self->ctx_weakreflist != NULL) {
470 PyObject_ClearWeakRefs((PyObject*)self);
471 }
472 (void)context_tp_clear(self);
473
474 #ifdef WITH_FREELISTS
475 struct _Py_context_freelist *context_freelist = get_context_freelist();
476 if (context_freelist->numfree >= 0 && context_freelist->numfree < PyContext_MAXFREELIST) {
477 context_freelist->numfree++;
478 self->ctx_weakreflist = (PyObject *)context_freelist->items;
479 context_freelist->items = self;
480 OBJECT_STAT_INC(to_freelist);
481 }
482 else
483 #endif
484 {
485 Py_TYPE(self)->tp_free(self);
486 }
487 }
488
489 static PyObject *
context_tp_iter(PyContext * self)490 context_tp_iter(PyContext *self)
491 {
492 return _PyHamt_NewIterKeys(self->ctx_vars);
493 }
494
495 static PyObject *
context_tp_richcompare(PyObject * v,PyObject * w,int op)496 context_tp_richcompare(PyObject *v, PyObject *w, int op)
497 {
498 if (!PyContext_CheckExact(v) || !PyContext_CheckExact(w) ||
499 (op != Py_EQ && op != Py_NE))
500 {
501 Py_RETURN_NOTIMPLEMENTED;
502 }
503
504 int res = _PyHamt_Eq(
505 ((PyContext *)v)->ctx_vars, ((PyContext *)w)->ctx_vars);
506 if (res < 0) {
507 return NULL;
508 }
509
510 if (op == Py_NE) {
511 res = !res;
512 }
513
514 if (res) {
515 Py_RETURN_TRUE;
516 }
517 else {
518 Py_RETURN_FALSE;
519 }
520 }
521
522 static Py_ssize_t
context_tp_len(PyContext * self)523 context_tp_len(PyContext *self)
524 {
525 return _PyHamt_Len(self->ctx_vars);
526 }
527
528 static PyObject *
context_tp_subscript(PyContext * self,PyObject * key)529 context_tp_subscript(PyContext *self, PyObject *key)
530 {
531 if (context_check_key_type(key)) {
532 return NULL;
533 }
534 PyObject *val = NULL;
535 int found = _PyHamt_Find(self->ctx_vars, key, &val);
536 if (found < 0) {
537 return NULL;
538 }
539 if (found == 0) {
540 PyErr_SetObject(PyExc_KeyError, key);
541 return NULL;
542 }
543 return Py_NewRef(val);
544 }
545
546 static int
context_tp_contains(PyContext * self,PyObject * key)547 context_tp_contains(PyContext *self, PyObject *key)
548 {
549 if (context_check_key_type(key)) {
550 return -1;
551 }
552 PyObject *val = NULL;
553 return _PyHamt_Find(self->ctx_vars, key, &val);
554 }
555
556
557 /*[clinic input]
558 _contextvars.Context.get
559 key: object
560 default: object = None
561 /
562
563 Return the value for `key` if `key` has the value in the context object.
564
565 If `key` does not exist, return `default`. If `default` is not given,
566 return None.
567 [clinic start generated code]*/
568
569 static PyObject *
_contextvars_Context_get_impl(PyContext * self,PyObject * key,PyObject * default_value)570 _contextvars_Context_get_impl(PyContext *self, PyObject *key,
571 PyObject *default_value)
572 /*[clinic end generated code: output=0c54aa7664268189 input=c8eeb81505023995]*/
573 {
574 if (context_check_key_type(key)) {
575 return NULL;
576 }
577
578 PyObject *val = NULL;
579 int found = _PyHamt_Find(self->ctx_vars, key, &val);
580 if (found < 0) {
581 return NULL;
582 }
583 if (found == 0) {
584 return Py_NewRef(default_value);
585 }
586 return Py_NewRef(val);
587 }
588
589
590 /*[clinic input]
591 _contextvars.Context.items
592
593 Return all variables and their values in the context object.
594
595 The result is returned as a list of 2-tuples (variable, value).
596 [clinic start generated code]*/
597
598 static PyObject *
_contextvars_Context_items_impl(PyContext * self)599 _contextvars_Context_items_impl(PyContext *self)
600 /*[clinic end generated code: output=fa1655c8a08502af input=00db64ae379f9f42]*/
601 {
602 return _PyHamt_NewIterItems(self->ctx_vars);
603 }
604
605
606 /*[clinic input]
607 _contextvars.Context.keys
608
609 Return a list of all variables in the context object.
610 [clinic start generated code]*/
611
612 static PyObject *
_contextvars_Context_keys_impl(PyContext * self)613 _contextvars_Context_keys_impl(PyContext *self)
614 /*[clinic end generated code: output=177227c6b63ec0e2 input=114b53aebca3449c]*/
615 {
616 return _PyHamt_NewIterKeys(self->ctx_vars);
617 }
618
619
620 /*[clinic input]
621 _contextvars.Context.values
622
623 Return a list of all variables' values in the context object.
624 [clinic start generated code]*/
625
626 static PyObject *
_contextvars_Context_values_impl(PyContext * self)627 _contextvars_Context_values_impl(PyContext *self)
628 /*[clinic end generated code: output=d286dabfc8db6dde input=ce8075d04a6ea526]*/
629 {
630 return _PyHamt_NewIterValues(self->ctx_vars);
631 }
632
633
634 /*[clinic input]
635 _contextvars.Context.copy
636
637 Return a shallow copy of the context object.
638 [clinic start generated code]*/
639
640 static PyObject *
_contextvars_Context_copy_impl(PyContext * self)641 _contextvars_Context_copy_impl(PyContext *self)
642 /*[clinic end generated code: output=30ba8896c4707a15 input=ebafdbdd9c72d592]*/
643 {
644 return (PyObject *)context_new_from_vars(self->ctx_vars);
645 }
646
647
648 static PyObject *
context_run(PyContext * self,PyObject * const * args,Py_ssize_t nargs,PyObject * kwnames)649 context_run(PyContext *self, PyObject *const *args,
650 Py_ssize_t nargs, PyObject *kwnames)
651 {
652 PyThreadState *ts = _PyThreadState_GET();
653
654 if (nargs < 1) {
655 _PyErr_SetString(ts, PyExc_TypeError,
656 "run() missing 1 required positional argument");
657 return NULL;
658 }
659
660 if (_PyContext_Enter(ts, (PyObject *)self)) {
661 return NULL;
662 }
663
664 PyObject *call_result = _PyObject_VectorcallTstate(
665 ts, args[0], args + 1, nargs - 1, kwnames);
666
667 if (_PyContext_Exit(ts, (PyObject *)self)) {
668 Py_XDECREF(call_result);
669 return NULL;
670 }
671
672 return call_result;
673 }
674
675
676 static PyMethodDef PyContext_methods[] = {
677 _CONTEXTVARS_CONTEXT_GET_METHODDEF
678 _CONTEXTVARS_CONTEXT_ITEMS_METHODDEF
679 _CONTEXTVARS_CONTEXT_KEYS_METHODDEF
680 _CONTEXTVARS_CONTEXT_VALUES_METHODDEF
681 _CONTEXTVARS_CONTEXT_COPY_METHODDEF
682 {"run", _PyCFunction_CAST(context_run), METH_FASTCALL | METH_KEYWORDS, NULL},
683 {NULL, NULL}
684 };
685
686 static PySequenceMethods PyContext_as_sequence = {
687 0, /* sq_length */
688 0, /* sq_concat */
689 0, /* sq_repeat */
690 0, /* sq_item */
691 0, /* sq_slice */
692 0, /* sq_ass_item */
693 0, /* sq_ass_slice */
694 (objobjproc)context_tp_contains, /* sq_contains */
695 0, /* sq_inplace_concat */
696 0, /* sq_inplace_repeat */
697 };
698
699 static PyMappingMethods PyContext_as_mapping = {
700 (lenfunc)context_tp_len, /* mp_length */
701 (binaryfunc)context_tp_subscript, /* mp_subscript */
702 };
703
704 PyTypeObject PyContext_Type = {
705 PyVarObject_HEAD_INIT(&PyType_Type, 0)
706 "_contextvars.Context",
707 sizeof(PyContext),
708 .tp_methods = PyContext_methods,
709 .tp_as_mapping = &PyContext_as_mapping,
710 .tp_as_sequence = &PyContext_as_sequence,
711 .tp_iter = (getiterfunc)context_tp_iter,
712 .tp_dealloc = (destructor)context_tp_dealloc,
713 .tp_getattro = PyObject_GenericGetAttr,
714 .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
715 .tp_richcompare = context_tp_richcompare,
716 .tp_traverse = (traverseproc)context_tp_traverse,
717 .tp_clear = (inquiry)context_tp_clear,
718 .tp_new = context_tp_new,
719 .tp_weaklistoffset = offsetof(PyContext, ctx_weakreflist),
720 .tp_hash = PyObject_HashNotImplemented,
721 };
722
723
724 /////////////////////////// ContextVar
725
726
727 static int
contextvar_set(PyContextVar * var,PyObject * val)728 contextvar_set(PyContextVar *var, PyObject *val)
729 {
730 #ifndef Py_GIL_DISABLED
731 var->var_cached = NULL;
732 PyThreadState *ts = _PyThreadState_GET();
733 #endif
734
735 PyContext *ctx = context_get();
736 if (ctx == NULL) {
737 return -1;
738 }
739
740 PyHamtObject *new_vars = _PyHamt_Assoc(
741 ctx->ctx_vars, (PyObject *)var, val);
742 if (new_vars == NULL) {
743 return -1;
744 }
745
746 Py_SETREF(ctx->ctx_vars, new_vars);
747
748 #ifndef Py_GIL_DISABLED
749 var->var_cached = val; /* borrow */
750 var->var_cached_tsid = ts->id;
751 var->var_cached_tsver = ts->context_ver;
752 #endif
753 return 0;
754 }
755
756 static int
contextvar_del(PyContextVar * var)757 contextvar_del(PyContextVar *var)
758 {
759 #ifndef Py_GIL_DISABLED
760 var->var_cached = NULL;
761 #endif
762
763 PyContext *ctx = context_get();
764 if (ctx == NULL) {
765 return -1;
766 }
767
768 PyHamtObject *vars = ctx->ctx_vars;
769 PyHamtObject *new_vars = _PyHamt_Without(vars, (PyObject *)var);
770 if (new_vars == NULL) {
771 return -1;
772 }
773
774 if (vars == new_vars) {
775 Py_DECREF(new_vars);
776 PyErr_SetObject(PyExc_LookupError, (PyObject *)var);
777 return -1;
778 }
779
780 Py_SETREF(ctx->ctx_vars, new_vars);
781 return 0;
782 }
783
784 static Py_hash_t
contextvar_generate_hash(void * addr,PyObject * name)785 contextvar_generate_hash(void *addr, PyObject *name)
786 {
787 /* Take hash of `name` and XOR it with the object's addr.
788
789 The structure of the tree is encoded in objects' hashes, which
790 means that sufficiently similar hashes would result in tall trees
791 with many Collision nodes. Which would, in turn, result in slower
792 get and set operations.
793
794 The XORing helps to ensure that:
795
796 (1) sequentially allocated ContextVar objects have
797 different hashes;
798
799 (2) context variables with equal names have
800 different hashes.
801 */
802
803 Py_hash_t name_hash = PyObject_Hash(name);
804 if (name_hash == -1) {
805 return -1;
806 }
807
808 Py_hash_t res = _Py_HashPointer(addr) ^ name_hash;
809 return res == -1 ? -2 : res;
810 }
811
812 static PyContextVar *
contextvar_new(PyObject * name,PyObject * def)813 contextvar_new(PyObject *name, PyObject *def)
814 {
815 if (!PyUnicode_Check(name)) {
816 PyErr_SetString(PyExc_TypeError,
817 "context variable name must be a str");
818 return NULL;
819 }
820
821 PyContextVar *var = PyObject_GC_New(PyContextVar, &PyContextVar_Type);
822 if (var == NULL) {
823 return NULL;
824 }
825
826 var->var_hash = contextvar_generate_hash(var, name);
827 if (var->var_hash == -1) {
828 Py_DECREF(var);
829 return NULL;
830 }
831
832 var->var_name = Py_NewRef(name);
833
834 var->var_default = Py_XNewRef(def);
835
836 #ifndef Py_GIL_DISABLED
837 var->var_cached = NULL;
838 var->var_cached_tsid = 0;
839 var->var_cached_tsver = 0;
840 #endif
841
842 if (_PyObject_GC_MAY_BE_TRACKED(name) ||
843 (def != NULL && _PyObject_GC_MAY_BE_TRACKED(def)))
844 {
845 PyObject_GC_Track(var);
846 }
847 return var;
848 }
849
850
851 /*[clinic input]
852 class _contextvars.ContextVar "PyContextVar *" "&PyContextVar_Type"
853 [clinic start generated code]*/
854 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=445da935fa8883c3]*/
855
856
857 static PyObject *
contextvar_tp_new(PyTypeObject * type,PyObject * args,PyObject * kwds)858 contextvar_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
859 {
860 static char *kwlist[] = {"", "default", NULL};
861 PyObject *name;
862 PyObject *def = NULL;
863
864 if (!PyArg_ParseTupleAndKeywords(
865 args, kwds, "O|$O:ContextVar", kwlist, &name, &def))
866 {
867 return NULL;
868 }
869
870 return (PyObject *)contextvar_new(name, def);
871 }
872
873 static int
contextvar_tp_clear(PyContextVar * self)874 contextvar_tp_clear(PyContextVar *self)
875 {
876 Py_CLEAR(self->var_name);
877 Py_CLEAR(self->var_default);
878 #ifndef Py_GIL_DISABLED
879 self->var_cached = NULL;
880 self->var_cached_tsid = 0;
881 self->var_cached_tsver = 0;
882 #endif
883 return 0;
884 }
885
886 static int
contextvar_tp_traverse(PyContextVar * self,visitproc visit,void * arg)887 contextvar_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
894 static void
contextvar_tp_dealloc(PyContextVar * self)895 contextvar_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
902 static Py_hash_t
contextvar_tp_hash(PyContextVar * self)903 contextvar_tp_hash(PyContextVar *self)
904 {
905 return self->var_hash;
906 }
907
908 static PyObject *
contextvar_tp_repr(PyContextVar * self)909 contextvar_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
959 error:
960 _PyUnicodeWriter_Dealloc(&writer);
961 return NULL;
962 }
963
964
965 /*[clinic input]
966 _contextvars.ContextVar.get
967 default: object = NULL
968 /
969
970 Return a value for the context variable for the current context.
971
972 If 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
979 static PyObject *
_contextvars_ContextVar_get_impl(PyContextVar * self,PyObject * default_value)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
1007 Call to set a new value for the context variable in the current context.
1008
1009 The required value argument is the new value for the context variable.
1010
1011 Returns a Token object that can be used to restore the variable to its previous
1012 value via the `ContextVar.reset()` method.
1013 [clinic start generated code]*/
1014
1015 static PyObject *
_contextvars_ContextVar_set(PyContextVar * self,PyObject * value)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
1027 Reset the context variable.
1028
1029 The variable is reset to the value it had before the `ContextVar.set()` that
1030 created the token was used.
1031 [clinic start generated code]*/
1032
1033 static PyObject *
_contextvars_ContextVar_reset(PyContextVar * self,PyObject * token)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
1051 static PyMemberDef PyContextVar_members[] = {
1052 {"name", _Py_T_OBJECT, offsetof(PyContextVar, var_name), Py_READONLY},
1053 {NULL}
1054 };
1055
1056 static 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
1065 PyTypeObject 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
1085 static PyObject * get_token_missing(void);
1086
1087
1088 /*[clinic input]
1089 class _contextvars.Token "PyContextToken *" "&PyContextToken_Type"
1090 [clinic start generated code]*/
1091 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=338a5e2db13d3f5b]*/
1092
1093
1094 static PyObject *
token_tp_new(PyTypeObject * type,PyObject * args,PyObject * kwds)1095 token_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
1102 static int
token_tp_clear(PyContextToken * self)1103 token_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
1111 static int
token_tp_traverse(PyContextToken * self,visitproc visit,void * arg)1112 token_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
1120 static void
token_tp_dealloc(PyContextToken * self)1121 token_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
1128 static PyObject *
token_tp_repr(PyContextToken * self)1129 token_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
1171 error:
1172 _PyUnicodeWriter_Dealloc(&writer);
1173 return NULL;
1174 }
1175
1176 static PyObject *
token_get_var(PyContextToken * self,void * Py_UNUSED (ignored))1177 token_get_var(PyContextToken *self, void *Py_UNUSED(ignored))
1178 {
1179 return Py_NewRef(self->tok_var);;
1180 }
1181
1182 static PyObject *
token_get_old_value(PyContextToken * self,void * Py_UNUSED (ignored))1183 token_get_old_value(PyContextToken *self, void *Py_UNUSED(ignored))
1184 {
1185 if (self->tok_oldval == NULL) {
1186 return get_token_missing();
1187 }
1188
1189 return Py_NewRef(self->tok_oldval);
1190 }
1191
1192 static PyGetSetDef PyContextTokenType_getsetlist[] = {
1193 {"var", (getter)token_get_var, NULL, NULL},
1194 {"old_value", (getter)token_get_old_value, NULL, NULL},
1195 {NULL}
1196 };
1197
1198 static PyMethodDef PyContextTokenType_methods[] = {
1199 {"__class_getitem__", Py_GenericAlias,
1200 METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
1201 {NULL}
1202 };
1203
1204 PyTypeObject PyContextToken_Type = {
1205 PyVarObject_HEAD_INIT(&PyType_Type, 0)
1206 "_contextvars.Token",
1207 sizeof(PyContextToken),
1208 .tp_methods = PyContextTokenType_methods,
1209 .tp_getset = PyContextTokenType_getsetlist,
1210 .tp_dealloc = (destructor)token_tp_dealloc,
1211 .tp_getattro = PyObject_GenericGetAttr,
1212 .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
1213 .tp_traverse = (traverseproc)token_tp_traverse,
1214 .tp_clear = (inquiry)token_tp_clear,
1215 .tp_new = token_tp_new,
1216 .tp_free = PyObject_GC_Del,
1217 .tp_hash = PyObject_HashNotImplemented,
1218 .tp_repr = (reprfunc)token_tp_repr,
1219 };
1220
1221 static PyContextToken *
token_new(PyContext * ctx,PyContextVar * var,PyObject * val)1222 token_new(PyContext *ctx, PyContextVar *var, PyObject *val)
1223 {
1224 PyContextToken *tok = PyObject_GC_New(PyContextToken, &PyContextToken_Type);
1225 if (tok == NULL) {
1226 return NULL;
1227 }
1228
1229 tok->tok_ctx = (PyContext*)Py_NewRef(ctx);
1230
1231 tok->tok_var = (PyContextVar*)Py_NewRef(var);
1232
1233 tok->tok_oldval = Py_XNewRef(val);
1234
1235 tok->tok_used = 0;
1236
1237 PyObject_GC_Track(tok);
1238 return tok;
1239 }
1240
1241
1242 /////////////////////////// Token.MISSING
1243
1244
1245 static PyObject *
context_token_missing_tp_repr(PyObject * self)1246 context_token_missing_tp_repr(PyObject *self)
1247 {
1248 return PyUnicode_FromString("<Token.MISSING>");
1249 }
1250
1251 static void
context_token_missing_tp_dealloc(_PyContextTokenMissing * Py_UNUSED (self))1252 context_token_missing_tp_dealloc(_PyContextTokenMissing *Py_UNUSED(self))
1253 {
1254 #ifdef Py_DEBUG
1255 /* The singleton is statically allocated. */
1256 _Py_FatalRefcountError("deallocating the token missing singleton");
1257 #else
1258 return;
1259 #endif
1260 }
1261
1262
1263 PyTypeObject _PyContextTokenMissing_Type = {
1264 PyVarObject_HEAD_INIT(&PyType_Type, 0)
1265 "Token.MISSING",
1266 sizeof(_PyContextTokenMissing),
1267 .tp_dealloc = (destructor)context_token_missing_tp_dealloc,
1268 .tp_getattro = PyObject_GenericGetAttr,
1269 .tp_flags = Py_TPFLAGS_DEFAULT,
1270 .tp_repr = context_token_missing_tp_repr,
1271 };
1272
1273
1274 static PyObject *
get_token_missing(void)1275 get_token_missing(void)
1276 {
1277 return (PyObject *)&_Py_SINGLETON(context_token_missing);
1278 }
1279
1280
1281 ///////////////////////////
1282
1283
1284 void
_PyContext_ClearFreeList(struct _Py_object_freelists * freelists,int is_finalization)1285 _PyContext_ClearFreeList(struct _Py_object_freelists *freelists, int is_finalization)
1286 {
1287 #ifdef WITH_FREELISTS
1288 struct _Py_context_freelist *state = &freelists->contexts;
1289 for (; state->numfree > 0; state->numfree--) {
1290 PyContext *ctx = state->items;
1291 state->items = (PyContext *)ctx->ctx_weakreflist;
1292 ctx->ctx_weakreflist = NULL;
1293 PyObject_GC_Del(ctx);
1294 }
1295 if (is_finalization) {
1296 state->numfree = -1;
1297 }
1298 #endif
1299 }
1300
1301
1302 PyStatus
_PyContext_Init(PyInterpreterState * interp)1303 _PyContext_Init(PyInterpreterState *interp)
1304 {
1305 if (!_Py_IsMainInterpreter(interp)) {
1306 return _PyStatus_OK();
1307 }
1308
1309 PyObject *missing = get_token_missing();
1310 if (PyDict_SetItemString(
1311 _PyType_GetDict(&PyContextToken_Type), "MISSING", missing))
1312 {
1313 Py_DECREF(missing);
1314 return _PyStatus_ERR("can't init context types");
1315 }
1316 Py_DECREF(missing);
1317
1318 return _PyStatus_OK();
1319 }
1320