Line data Source code
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]
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 : #if PyContext_MAXFREELIST > 0
68 : static struct _Py_context_state *
69 74778 : get_context_state(void)
70 : {
71 74778 : PyInterpreterState *interp = _PyInterpreterState_GET();
72 74778 : return &interp->context;
73 : }
74 : #endif
75 :
76 :
77 : PyObject *
78 26 : _PyContext_NewHamtForTests(void)
79 : {
80 26 : return (PyObject *)_PyHamt_New();
81 : }
82 :
83 :
84 : PyObject *
85 21 : PyContext_New(void)
86 : {
87 21 : return (PyObject *)context_new_empty();
88 : }
89 :
90 :
91 : PyObject *
92 0 : PyContext_Copy(PyObject * octx)
93 : {
94 0 : ENSURE_Context(octx, NULL)
95 0 : PyContext *ctx = (PyContext *)octx;
96 0 : return (PyObject *)context_new_from_vars(ctx->ctx_vars);
97 : }
98 :
99 :
100 : PyObject *
101 37097 : PyContext_CopyCurrent(void)
102 : {
103 37097 : PyContext *ctx = context_get();
104 37097 : if (ctx == NULL) {
105 0 : return NULL;
106 : }
107 :
108 37097 : return (PyObject *)context_new_from_vars(ctx->ctx_vars);
109 : }
110 :
111 :
112 : static int
113 53013 : _PyContext_Enter(PyThreadState *ts, PyObject *octx)
114 : {
115 53013 : ENSURE_Context(octx, -1)
116 53013 : PyContext *ctx = (PyContext *)octx;
117 :
118 53013 : if (ctx->ctx_entered) {
119 1 : _PyErr_Format(ts, PyExc_RuntimeError,
120 : "cannot enter context: %R is already entered", ctx);
121 1 : return -1;
122 : }
123 :
124 53012 : ctx->ctx_prev = (PyContext *)ts->context; /* borrow */
125 53012 : ctx->ctx_entered = 1;
126 :
127 53012 : Py_INCREF(ctx);
128 53012 : ts->context = (PyObject *)ctx;
129 53012 : ts->context_ver++;
130 :
131 53012 : return 0;
132 : }
133 :
134 :
135 : int
136 0 : PyContext_Enter(PyObject *octx)
137 : {
138 0 : PyThreadState *ts = _PyThreadState_GET();
139 0 : assert(ts != NULL);
140 0 : return _PyContext_Enter(ts, octx);
141 : }
142 :
143 :
144 : static int
145 53012 : _PyContext_Exit(PyThreadState *ts, PyObject *octx)
146 : {
147 53012 : ENSURE_Context(octx, -1)
148 53012 : PyContext *ctx = (PyContext *)octx;
149 :
150 53012 : if (!ctx->ctx_entered) {
151 0 : PyErr_Format(PyExc_RuntimeError,
152 : "cannot exit context: %R has not been entered", ctx);
153 0 : return -1;
154 : }
155 :
156 53012 : if (ts->context != (PyObject *)ctx) {
157 : /* Can only happen if someone misuses the C API */
158 0 : PyErr_SetString(PyExc_RuntimeError,
159 : "cannot exit context: thread state references "
160 : "a different context object");
161 0 : return -1;
162 : }
163 :
164 53012 : Py_SETREF(ts->context, (PyObject *)ctx->ctx_prev);
165 53012 : ts->context_ver++;
166 :
167 53012 : ctx->ctx_prev = NULL;
168 53012 : ctx->ctx_entered = 0;
169 :
170 53012 : return 0;
171 : }
172 :
173 : int
174 0 : PyContext_Exit(PyObject *octx)
175 : {
176 0 : PyThreadState *ts = _PyThreadState_GET();
177 0 : assert(ts != NULL);
178 0 : return _PyContext_Exit(ts, octx);
179 : }
180 :
181 :
182 : PyObject *
183 50 : PyContextVar_New(const char *name, PyObject *def)
184 : {
185 50 : PyObject *pyname = PyUnicode_FromString(name);
186 50 : if (pyname == NULL) {
187 0 : return NULL;
188 : }
189 50 : PyContextVar *var = contextvar_new(pyname, def);
190 50 : Py_DECREF(pyname);
191 50 : return (PyObject *)var;
192 : }
193 :
194 :
195 : int
196 280325 : PyContextVar_Get(PyObject *ovar, PyObject *def, PyObject **val)
197 : {
198 280325 : ENSURE_ContextVar(ovar, -1)
199 280325 : PyContextVar *var = (PyContextVar *)ovar;
200 :
201 280325 : PyThreadState *ts = _PyThreadState_GET();
202 280325 : assert(ts != NULL);
203 280325 : if (ts->context == NULL) {
204 22 : goto not_found;
205 : }
206 :
207 280303 : if (var->var_cached != NULL &&
208 280275 : var->var_cached_tsid == ts->id &&
209 280174 : var->var_cached_tsver == ts->context_ver)
210 : {
211 272054 : *val = var->var_cached;
212 272054 : goto found;
213 : }
214 :
215 8249 : assert(PyContext_CheckExact(ts->context));
216 8249 : PyHamtObject *vars = ((PyContext *)ts->context)->ctx_vars;
217 :
218 8249 : PyObject *found = NULL;
219 8249 : int res = _PyHamt_Find(vars, (PyObject*)var, &found);
220 8249 : if (res < 0) {
221 0 : goto error;
222 : }
223 8249 : if (res == 1) {
224 8192 : assert(found != NULL);
225 8192 : var->var_cached = found; /* borrow */
226 8192 : var->var_cached_tsid = ts->id;
227 8192 : var->var_cached_tsver = ts->context_ver;
228 :
229 8192 : *val = found;
230 8192 : goto found;
231 : }
232 :
233 57 : not_found:
234 79 : if (def == NULL) {
235 73 : if (var->var_default != NULL) {
236 45 : *val = var->var_default;
237 45 : goto found;
238 : }
239 :
240 28 : *val = NULL;
241 28 : goto found;
242 : }
243 : else {
244 6 : *val = def;
245 6 : goto found;
246 : }
247 :
248 280325 : found:
249 280325 : Py_XINCREF(*val);
250 280325 : return 0;
251 :
252 0 : error:
253 0 : *val = NULL;
254 0 : return -1;
255 : }
256 :
257 :
258 : PyObject *
259 8717 : PyContextVar_Set(PyObject *ovar, PyObject *val)
260 : {
261 8717 : ENSURE_ContextVar(ovar, NULL)
262 8717 : PyContextVar *var = (PyContextVar *)ovar;
263 :
264 8717 : if (!PyContextVar_CheckExact(var)) {
265 0 : PyErr_SetString(
266 : PyExc_TypeError, "an instance of ContextVar was expected");
267 0 : return NULL;
268 : }
269 :
270 8717 : PyContext *ctx = context_get();
271 8717 : if (ctx == NULL) {
272 0 : return NULL;
273 : }
274 :
275 8717 : PyObject *old_val = NULL;
276 8717 : int found = _PyHamt_Find(ctx->ctx_vars, (PyObject *)var, &old_val);
277 8717 : if (found < 0) {
278 0 : return NULL;
279 : }
280 :
281 8717 : Py_XINCREF(old_val);
282 8717 : PyContextToken *tok = token_new(ctx, var, old_val);
283 8717 : Py_XDECREF(old_val);
284 :
285 8717 : if (contextvar_set(var, val)) {
286 0 : Py_DECREF(tok);
287 0 : return NULL;
288 : }
289 :
290 8717 : return (PyObject *)tok;
291 : }
292 :
293 :
294 : int
295 8 : PyContextVar_Reset(PyObject *ovar, PyObject *otok)
296 : {
297 8 : ENSURE_ContextVar(ovar, -1)
298 8 : ENSURE_ContextToken(otok, -1)
299 8 : PyContextVar *var = (PyContextVar *)ovar;
300 8 : PyContextToken *tok = (PyContextToken *)otok;
301 :
302 8 : if (tok->tok_used) {
303 2 : PyErr_Format(PyExc_RuntimeError,
304 : "%R has already been used once", tok);
305 2 : return -1;
306 : }
307 :
308 6 : if (var != tok->tok_var) {
309 1 : PyErr_Format(PyExc_ValueError,
310 : "%R was created by a different ContextVar", tok);
311 1 : return -1;
312 : }
313 :
314 5 : PyContext *ctx = context_get();
315 5 : if (ctx != tok->tok_ctx) {
316 1 : PyErr_Format(PyExc_ValueError,
317 : "%R was created in a different Context", tok);
318 1 : return -1;
319 : }
320 :
321 4 : tok->tok_used = 1;
322 :
323 4 : if (tok->tok_oldval == NULL) {
324 3 : return contextvar_del(var);
325 : }
326 : else {
327 1 : return contextvar_set(var, tok->tok_oldval);
328 : }
329 : }
330 :
331 :
332 : /////////////////////////// PyContext
333 :
334 : /*[clinic input]
335 : class _contextvars.Context "PyContext *" "&PyContext_Type"
336 : [clinic start generated code]*/
337 : /*[clinic end generated code: output=da39a3ee5e6b4b0d input=bdf87f8e0cb580e8]*/
338 :
339 :
340 : static inline PyContext *
341 37389 : _context_alloc(void)
342 : {
343 : PyContext *ctx;
344 : #if PyContext_MAXFREELIST > 0
345 37389 : struct _Py_context_state *state = get_context_state();
346 : #ifdef Py_DEBUG
347 : // _context_alloc() must not be called after _PyContext_Fini()
348 37389 : assert(state->numfree != -1);
349 : #endif
350 37389 : if (state->numfree) {
351 28728 : state->numfree--;
352 28728 : ctx = state->freelist;
353 28728 : state->freelist = (PyContext *)ctx->ctx_weakreflist;
354 : OBJECT_STAT_INC(from_freelist);
355 28728 : ctx->ctx_weakreflist = NULL;
356 28728 : _Py_NewReference((PyObject *)ctx);
357 : }
358 : else
359 : #endif
360 : {
361 8661 : ctx = PyObject_GC_New(PyContext, &PyContext_Type);
362 8661 : if (ctx == NULL) {
363 0 : return NULL;
364 : }
365 : }
366 :
367 37389 : ctx->ctx_vars = NULL;
368 37389 : ctx->ctx_prev = NULL;
369 37389 : ctx->ctx_entered = 0;
370 37389 : ctx->ctx_weakreflist = NULL;
371 :
372 37389 : return ctx;
373 : }
374 :
375 :
376 : static PyContext *
377 291 : context_new_empty(void)
378 : {
379 291 : PyContext *ctx = _context_alloc();
380 291 : if (ctx == NULL) {
381 0 : return NULL;
382 : }
383 :
384 291 : ctx->ctx_vars = _PyHamt_New();
385 291 : if (ctx->ctx_vars == NULL) {
386 0 : Py_DECREF(ctx);
387 0 : return NULL;
388 : }
389 :
390 291 : _PyObject_GC_TRACK(ctx);
391 291 : return ctx;
392 : }
393 :
394 :
395 : static PyContext *
396 37098 : context_new_from_vars(PyHamtObject *vars)
397 : {
398 37098 : PyContext *ctx = _context_alloc();
399 37098 : if (ctx == NULL) {
400 0 : return NULL;
401 : }
402 :
403 37098 : Py_INCREF(vars);
404 37098 : ctx->ctx_vars = vars;
405 :
406 37098 : _PyObject_GC_TRACK(ctx);
407 37098 : return ctx;
408 : }
409 :
410 :
411 : static inline PyContext *
412 54540 : context_get(void)
413 : {
414 54540 : PyThreadState *ts = _PyThreadState_GET();
415 54540 : assert(ts != NULL);
416 54540 : PyContext *current_ctx = (PyContext *)ts->context;
417 54540 : if (current_ctx == NULL) {
418 270 : current_ctx = context_new_empty();
419 270 : if (current_ctx == NULL) {
420 0 : return NULL;
421 : }
422 270 : ts->context = (PyObject *)current_ctx;
423 : }
424 54540 : return current_ctx;
425 : }
426 :
427 : static int
428 102 : context_check_key_type(PyObject *key)
429 : {
430 102 : if (!PyContextVar_CheckExact(key)) {
431 : // abort();
432 3 : PyErr_Format(PyExc_TypeError,
433 : "a ContextVar key was expected, got %R", key);
434 3 : return -1;
435 : }
436 99 : return 0;
437 : }
438 :
439 : static PyObject *
440 24 : context_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
441 : {
442 24 : if (PyTuple_Size(args) || (kwds != NULL && PyDict_Size(kwds))) {
443 3 : PyErr_SetString(
444 : PyExc_TypeError, "Context() does not accept any arguments");
445 3 : return NULL;
446 : }
447 21 : return PyContext_New();
448 : }
449 :
450 : static int
451 37402 : context_tp_clear(PyContext *self)
452 : {
453 37402 : Py_CLEAR(self->ctx_prev);
454 37402 : Py_CLEAR(self->ctx_vars);
455 37402 : return 0;
456 : }
457 :
458 : static int
459 505400 : context_tp_traverse(PyContext *self, visitproc visit, void *arg)
460 : {
461 505400 : Py_VISIT(self->ctx_prev);
462 505400 : Py_VISIT(self->ctx_vars);
463 505400 : return 0;
464 : }
465 :
466 : static void
467 37389 : context_tp_dealloc(PyContext *self)
468 : {
469 37389 : _PyObject_GC_UNTRACK(self);
470 :
471 37389 : if (self->ctx_weakreflist != NULL) {
472 0 : PyObject_ClearWeakRefs((PyObject*)self);
473 : }
474 37389 : (void)context_tp_clear(self);
475 :
476 : #if PyContext_MAXFREELIST > 0
477 37389 : struct _Py_context_state *state = get_context_state();
478 : #ifdef Py_DEBUG
479 : // _context_alloc() must not be called after _PyContext_Fini()
480 37389 : assert(state->numfree != -1);
481 : #endif
482 37389 : if (state->numfree < PyContext_MAXFREELIST) {
483 36994 : state->numfree++;
484 36994 : self->ctx_weakreflist = (PyObject *)state->freelist;
485 36994 : state->freelist = self;
486 : OBJECT_STAT_INC(to_freelist);
487 : }
488 : else
489 : #endif
490 : {
491 395 : Py_TYPE(self)->tp_free(self);
492 : }
493 37389 : }
494 :
495 : static PyObject *
496 2 : context_tp_iter(PyContext *self)
497 : {
498 2 : return _PyHamt_NewIterKeys(self->ctx_vars);
499 : }
500 :
501 : static PyObject *
502 1 : context_tp_richcompare(PyObject *v, PyObject *w, int op)
503 : {
504 1 : if (!PyContext_CheckExact(v) || !PyContext_CheckExact(w) ||
505 0 : (op != Py_EQ && op != Py_NE))
506 : {
507 0 : Py_RETURN_NOTIMPLEMENTED;
508 : }
509 :
510 1 : int res = _PyHamt_Eq(
511 : ((PyContext *)v)->ctx_vars, ((PyContext *)w)->ctx_vars);
512 1 : if (res < 0) {
513 0 : return NULL;
514 : }
515 :
516 1 : if (op == Py_NE) {
517 0 : res = !res;
518 : }
519 :
520 1 : if (res) {
521 1 : Py_RETURN_TRUE;
522 : }
523 : else {
524 0 : Py_RETURN_FALSE;
525 : }
526 : }
527 :
528 : static Py_ssize_t
529 5 : context_tp_len(PyContext *self)
530 : {
531 5 : return _PyHamt_Len(self->ctx_vars);
532 : }
533 :
534 : static PyObject *
535 13 : context_tp_subscript(PyContext *self, PyObject *key)
536 : {
537 13 : if (context_check_key_type(key)) {
538 1 : return NULL;
539 : }
540 12 : PyObject *val = NULL;
541 12 : int found = _PyHamt_Find(self->ctx_vars, key, &val);
542 12 : if (found < 0) {
543 0 : return NULL;
544 : }
545 12 : if (found == 0) {
546 3 : PyErr_SetObject(PyExc_KeyError, key);
547 3 : return NULL;
548 : }
549 9 : Py_INCREF(val);
550 9 : return val;
551 : }
552 :
553 : static int
554 6 : context_tp_contains(PyContext *self, PyObject *key)
555 : {
556 6 : if (context_check_key_type(key)) {
557 1 : return -1;
558 : }
559 5 : PyObject *val = NULL;
560 5 : 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 :
570 : Return the value for `key` if `key` has the value in the context object.
571 :
572 : If `key` does not exist, return `default`. If `default` is not given,
573 : return None.
574 : [clinic start generated code]*/
575 :
576 : static PyObject *
577 83 : _contextvars_Context_get_impl(PyContext *self, PyObject *key,
578 : PyObject *default_value)
579 : /*[clinic end generated code: output=0c54aa7664268189 input=c8eeb81505023995]*/
580 : {
581 83 : if (context_check_key_type(key)) {
582 1 : return NULL;
583 : }
584 :
585 82 : PyObject *val = NULL;
586 82 : int found = _PyHamt_Find(self->ctx_vars, key, &val);
587 82 : if (found < 0) {
588 0 : return NULL;
589 : }
590 82 : if (found == 0) {
591 29 : Py_INCREF(default_value);
592 29 : return default_value;
593 : }
594 53 : Py_INCREF(val);
595 53 : return val;
596 : }
597 :
598 :
599 : /*[clinic input]
600 : _contextvars.Context.items
601 :
602 : Return all variables and their values in the context object.
603 :
604 : The result is returned as a list of 2-tuples (variable, value).
605 : [clinic start generated code]*/
606 :
607 : static PyObject *
608 1 : _contextvars_Context_items_impl(PyContext *self)
609 : /*[clinic end generated code: output=fa1655c8a08502af input=00db64ae379f9f42]*/
610 : {
611 1 : return _PyHamt_NewIterItems(self->ctx_vars);
612 : }
613 :
614 :
615 : /*[clinic input]
616 : _contextvars.Context.keys
617 :
618 : Return a list of all variables in the context object.
619 : [clinic start generated code]*/
620 :
621 : static PyObject *
622 4 : _contextvars_Context_keys_impl(PyContext *self)
623 : /*[clinic end generated code: output=177227c6b63ec0e2 input=114b53aebca3449c]*/
624 : {
625 4 : return _PyHamt_NewIterKeys(self->ctx_vars);
626 : }
627 :
628 :
629 : /*[clinic input]
630 : _contextvars.Context.values
631 :
632 : Return a list of all variables' values in the context object.
633 : [clinic start generated code]*/
634 :
635 : static PyObject *
636 1 : _contextvars_Context_values_impl(PyContext *self)
637 : /*[clinic end generated code: output=d286dabfc8db6dde input=ce8075d04a6ea526]*/
638 : {
639 1 : return _PyHamt_NewIterValues(self->ctx_vars);
640 : }
641 :
642 :
643 : /*[clinic input]
644 : _contextvars.Context.copy
645 :
646 : Return a shallow copy of the context object.
647 : [clinic start generated code]*/
648 :
649 : static PyObject *
650 1 : _contextvars_Context_copy_impl(PyContext *self)
651 : /*[clinic end generated code: output=30ba8896c4707a15 input=ebafdbdd9c72d592]*/
652 : {
653 1 : return (PyObject *)context_new_from_vars(self->ctx_vars);
654 : }
655 :
656 :
657 : static PyObject *
658 53014 : context_run(PyContext *self, PyObject *const *args,
659 : Py_ssize_t nargs, PyObject *kwnames)
660 : {
661 53014 : PyThreadState *ts = _PyThreadState_GET();
662 :
663 53014 : if (nargs < 1) {
664 1 : _PyErr_SetString(ts, PyExc_TypeError,
665 : "run() missing 1 required positional argument");
666 1 : return NULL;
667 : }
668 :
669 53013 : if (_PyContext_Enter(ts, (PyObject *)self)) {
670 1 : return NULL;
671 : }
672 :
673 53012 : PyObject *call_result = _PyObject_VectorcallTstate(
674 53012 : ts, args[0], args + 1, nargs - 1, kwnames);
675 :
676 53012 : if (_PyContext_Exit(ts, (PyObject *)self)) {
677 0 : return NULL;
678 : }
679 :
680 53012 : return call_result;
681 : }
682 :
683 :
684 : static 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 :
694 : static 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 :
707 : static PyMappingMethods PyContext_as_mapping = {
708 : (lenfunc)context_tp_len, /* mp_length */
709 : (binaryfunc)context_tp_subscript, /* mp_subscript */
710 : };
711 :
712 : PyTypeObject 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 :
735 : static int
736 8718 : contextvar_set(PyContextVar *var, PyObject *val)
737 : {
738 8718 : var->var_cached = NULL;
739 8718 : PyThreadState *ts = _PyThreadState_GET();
740 :
741 8718 : PyContext *ctx = context_get();
742 8718 : if (ctx == NULL) {
743 0 : return -1;
744 : }
745 :
746 8718 : PyHamtObject *new_vars = _PyHamt_Assoc(
747 : ctx->ctx_vars, (PyObject *)var, val);
748 8718 : if (new_vars == NULL) {
749 0 : return -1;
750 : }
751 :
752 8718 : Py_SETREF(ctx->ctx_vars, new_vars);
753 :
754 8718 : var->var_cached = val; /* borrow */
755 8718 : var->var_cached_tsid = ts->id;
756 8718 : var->var_cached_tsver = ts->context_ver;
757 8718 : return 0;
758 : }
759 :
760 : static int
761 3 : contextvar_del(PyContextVar *var)
762 : {
763 3 : var->var_cached = NULL;
764 :
765 3 : PyContext *ctx = context_get();
766 3 : if (ctx == NULL) {
767 0 : return -1;
768 : }
769 :
770 3 : PyHamtObject *vars = ctx->ctx_vars;
771 3 : PyHamtObject *new_vars = _PyHamt_Without(vars, (PyObject *)var);
772 3 : if (new_vars == NULL) {
773 0 : return -1;
774 : }
775 :
776 3 : if (vars == new_vars) {
777 0 : Py_DECREF(new_vars);
778 0 : PyErr_SetObject(PyExc_LookupError, (PyObject *)var);
779 0 : return -1;
780 : }
781 :
782 3 : Py_SETREF(ctx->ctx_vars, new_vars);
783 3 : return 0;
784 : }
785 :
786 : static Py_hash_t
787 119 : contextvar_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 119 : Py_hash_t name_hash = PyObject_Hash(name);
806 119 : if (name_hash == -1) {
807 0 : return -1;
808 : }
809 :
810 119 : Py_hash_t res = _Py_HashPointer(addr) ^ name_hash;
811 119 : return res == -1 ? -2 : res;
812 : }
813 :
814 : static PyContextVar *
815 120 : contextvar_new(PyObject *name, PyObject *def)
816 : {
817 120 : if (!PyUnicode_Check(name)) {
818 1 : PyErr_SetString(PyExc_TypeError,
819 : "context variable name must be a str");
820 1 : return NULL;
821 : }
822 :
823 119 : PyContextVar *var = PyObject_GC_New(PyContextVar, &PyContextVar_Type);
824 119 : if (var == NULL) {
825 0 : return NULL;
826 : }
827 :
828 119 : var->var_hash = contextvar_generate_hash(var, name);
829 119 : if (var->var_hash == -1) {
830 0 : Py_DECREF(var);
831 0 : return NULL;
832 : }
833 :
834 119 : Py_INCREF(name);
835 119 : var->var_name = name;
836 :
837 119 : Py_XINCREF(def);
838 119 : var->var_default = def;
839 :
840 119 : var->var_cached = NULL;
841 119 : var->var_cached_tsid = 0;
842 119 : var->var_cached_tsver = 0;
843 :
844 119 : if (_PyObject_GC_MAY_BE_TRACKED(name) ||
845 33 : (def != NULL && _PyObject_GC_MAY_BE_TRACKED(def)))
846 : {
847 1 : PyObject_GC_Track(var);
848 : }
849 119 : return var;
850 : }
851 :
852 :
853 : /*[clinic input]
854 : class _contextvars.ContextVar "PyContextVar *" "&PyContextVar_Type"
855 : [clinic start generated code]*/
856 : /*[clinic end generated code: output=da39a3ee5e6b4b0d input=445da935fa8883c3]*/
857 :
858 :
859 : static PyObject *
860 71 : contextvar_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
861 : {
862 : static char *kwlist[] = {"", "default", NULL};
863 : PyObject *name;
864 71 : PyObject *def = NULL;
865 :
866 71 : if (!PyArg_ParseTupleAndKeywords(
867 : args, kwds, "O|$O:ContextVar", kwlist, &name, &def))
868 : {
869 1 : return NULL;
870 : }
871 :
872 70 : return (PyObject *)contextvar_new(name, def);
873 : }
874 :
875 : static int
876 69 : contextvar_tp_clear(PyContextVar *self)
877 : {
878 69 : Py_CLEAR(self->var_name);
879 69 : Py_CLEAR(self->var_default);
880 69 : self->var_cached = NULL;
881 69 : self->var_cached_tsid = 0;
882 69 : self->var_cached_tsver = 0;
883 69 : return 0;
884 : }
885 :
886 : static int
887 3 : contextvar_tp_traverse(PyContextVar *self, visitproc visit, void *arg)
888 : {
889 3 : Py_VISIT(self->var_name);
890 3 : Py_VISIT(self->var_default);
891 3 : return 0;
892 : }
893 :
894 : static void
895 69 : contextvar_tp_dealloc(PyContextVar *self)
896 : {
897 69 : PyObject_GC_UnTrack(self);
898 69 : (void)contextvar_tp_clear(self);
899 69 : Py_TYPE(self)->tp_free(self);
900 69 : }
901 :
902 : static Py_hash_t
903 25712 : contextvar_tp_hash(PyContextVar *self)
904 : {
905 25712 : return self->var_hash;
906 : }
907 :
908 : static PyObject *
909 17 : contextvar_tp_repr(PyContextVar *self)
910 : {
911 : _PyUnicodeWriter writer;
912 :
913 17 : _PyUnicodeWriter_Init(&writer);
914 :
915 17 : if (_PyUnicodeWriter_WriteASCIIString(
916 : &writer, "<ContextVar name=", 17) < 0)
917 : {
918 0 : goto error;
919 : }
920 :
921 17 : PyObject *name = PyObject_Repr(self->var_name);
922 17 : if (name == NULL) {
923 0 : goto error;
924 : }
925 17 : if (_PyUnicodeWriter_WriteStr(&writer, name) < 0) {
926 0 : Py_DECREF(name);
927 0 : goto error;
928 : }
929 17 : Py_DECREF(name);
930 :
931 17 : if (self->var_default != NULL) {
932 13 : if (_PyUnicodeWriter_WriteASCIIString(&writer, " default=", 9) < 0) {
933 0 : goto error;
934 : }
935 :
936 13 : PyObject *def = PyObject_Repr(self->var_default);
937 13 : if (def == NULL) {
938 0 : goto error;
939 : }
940 13 : if (_PyUnicodeWriter_WriteStr(&writer, def) < 0) {
941 0 : Py_DECREF(def);
942 0 : goto error;
943 : }
944 13 : Py_DECREF(def);
945 : }
946 :
947 17 : PyObject *addr = PyUnicode_FromFormat(" at %p>", self);
948 17 : if (addr == NULL) {
949 0 : goto error;
950 : }
951 17 : if (_PyUnicodeWriter_WriteStr(&writer, addr) < 0) {
952 0 : Py_DECREF(addr);
953 0 : goto error;
954 : }
955 17 : Py_DECREF(addr);
956 :
957 17 : return _PyUnicodeWriter_Finish(&writer);
958 :
959 0 : error:
960 0 : _PyUnicodeWriter_Dealloc(&writer);
961 0 : 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 *
980 42112 : _contextvars_ContextVar_get_impl(PyContextVar *self, PyObject *default_value)
981 : /*[clinic end generated code: output=0746bd0aa2ced7bf input=30aa2ab9e433e401]*/
982 : {
983 42112 : if (!PyContextVar_CheckExact(self)) {
984 0 : PyErr_SetString(
985 : PyExc_TypeError, "an instance of ContextVar was expected");
986 0 : return NULL;
987 : }
988 :
989 : PyObject *val;
990 42112 : if (PyContextVar_Get((PyObject *)self, default_value, &val) < 0) {
991 0 : return NULL;
992 : }
993 :
994 42112 : if (val == NULL) {
995 4 : PyErr_SetObject(PyExc_LookupError, (PyObject *)self);
996 4 : return NULL;
997 : }
998 :
999 42108 : 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 *
1016 8383 : _contextvars_ContextVar_set(PyContextVar *self, PyObject *value)
1017 : /*[clinic end generated code: output=446ed5e820d6d60b input=c0a6887154227453]*/
1018 : {
1019 8383 : 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 *
1034 8 : _contextvars_ContextVar_reset(PyContextVar *self, PyObject *token)
1035 : /*[clinic end generated code: output=d4ee34d0742d62ee input=ebe2881e5af4ffda]*/
1036 : {
1037 8 : if (!PyContextToken_CheckExact(token)) {
1038 0 : PyErr_Format(PyExc_TypeError,
1039 : "expected an instance of Token, got %R", token);
1040 0 : return NULL;
1041 : }
1042 :
1043 8 : if (PyContextVar_Reset((PyObject *)self, token)) {
1044 4 : return NULL;
1045 : }
1046 :
1047 4 : Py_RETURN_NONE;
1048 : }
1049 :
1050 :
1051 : static PyMemberDef PyContextVar_members[] = {
1052 : {"name", T_OBJECT, offsetof(PyContextVar, var_name), 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 *
1095 0 : token_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1096 : {
1097 0 : PyErr_SetString(PyExc_RuntimeError,
1098 : "Tokens can only be created by ContextVars");
1099 0 : return NULL;
1100 : }
1101 :
1102 : static int
1103 8717 : token_tp_clear(PyContextToken *self)
1104 : {
1105 8717 : Py_CLEAR(self->tok_ctx);
1106 8717 : Py_CLEAR(self->tok_var);
1107 8717 : Py_CLEAR(self->tok_oldval);
1108 8717 : return 0;
1109 : }
1110 :
1111 : static int
1112 8717 : token_tp_traverse(PyContextToken *self, visitproc visit, void *arg)
1113 : {
1114 8717 : Py_VISIT(self->tok_ctx);
1115 8717 : Py_VISIT(self->tok_var);
1116 8717 : Py_VISIT(self->tok_oldval);
1117 8717 : return 0;
1118 : }
1119 :
1120 : static void
1121 8717 : token_tp_dealloc(PyContextToken *self)
1122 : {
1123 8717 : PyObject_GC_UnTrack(self);
1124 8717 : (void)token_tp_clear(self);
1125 8717 : Py_TYPE(self)->tp_free(self);
1126 8717 : }
1127 :
1128 : static PyObject *
1129 7 : token_tp_repr(PyContextToken *self)
1130 : {
1131 : _PyUnicodeWriter writer;
1132 :
1133 7 : _PyUnicodeWriter_Init(&writer);
1134 :
1135 7 : if (_PyUnicodeWriter_WriteASCIIString(&writer, "<Token", 6) < 0) {
1136 0 : goto error;
1137 : }
1138 :
1139 7 : if (self->tok_used) {
1140 3 : if (_PyUnicodeWriter_WriteASCIIString(&writer, " used", 5) < 0) {
1141 0 : goto error;
1142 : }
1143 : }
1144 :
1145 7 : if (_PyUnicodeWriter_WriteASCIIString(&writer, " var=", 5) < 0) {
1146 0 : goto error;
1147 : }
1148 :
1149 7 : PyObject *var = PyObject_Repr((PyObject *)self->tok_var);
1150 7 : if (var == NULL) {
1151 0 : goto error;
1152 : }
1153 7 : if (_PyUnicodeWriter_WriteStr(&writer, var) < 0) {
1154 0 : Py_DECREF(var);
1155 0 : goto error;
1156 : }
1157 7 : Py_DECREF(var);
1158 :
1159 7 : PyObject *addr = PyUnicode_FromFormat(" at %p>", self);
1160 7 : if (addr == NULL) {
1161 0 : goto error;
1162 : }
1163 7 : if (_PyUnicodeWriter_WriteStr(&writer, addr) < 0) {
1164 0 : Py_DECREF(addr);
1165 0 : goto error;
1166 : }
1167 7 : Py_DECREF(addr);
1168 :
1169 7 : return _PyUnicodeWriter_Finish(&writer);
1170 :
1171 0 : error:
1172 0 : _PyUnicodeWriter_Dealloc(&writer);
1173 0 : return NULL;
1174 : }
1175 :
1176 : static PyObject *
1177 1 : token_get_var(PyContextToken *self, void *Py_UNUSED(ignored))
1178 : {
1179 1 : Py_INCREF(self->tok_var);
1180 1 : return (PyObject *)self->tok_var;
1181 : }
1182 :
1183 : static PyObject *
1184 3 : token_get_old_value(PyContextToken *self, void *Py_UNUSED(ignored))
1185 : {
1186 3 : if (self->tok_oldval == NULL) {
1187 2 : return get_token_missing();
1188 : }
1189 :
1190 1 : Py_INCREF(self->tok_oldval);
1191 1 : return self->tok_oldval;
1192 : }
1193 :
1194 : static 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 :
1200 : static PyMethodDef PyContextTokenType_methods[] = {
1201 : {"__class_getitem__", Py_GenericAlias,
1202 : METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
1203 : {NULL}
1204 : };
1205 :
1206 : PyTypeObject 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 :
1223 : static PyContextToken *
1224 8717 : token_new(PyContext *ctx, PyContextVar *var, PyObject *val)
1225 : {
1226 8717 : PyContextToken *tok = PyObject_GC_New(PyContextToken, &PyContextToken_Type);
1227 8717 : if (tok == NULL) {
1228 0 : return NULL;
1229 : }
1230 :
1231 8717 : Py_INCREF(ctx);
1232 8717 : tok->tok_ctx = ctx;
1233 :
1234 8717 : Py_INCREF(var);
1235 8717 : tok->tok_var = var;
1236 :
1237 8717 : Py_XINCREF(val);
1238 8717 : tok->tok_oldval = val;
1239 :
1240 8717 : tok->tok_used = 0;
1241 :
1242 8717 : PyObject_GC_Track(tok);
1243 8717 : return tok;
1244 : }
1245 :
1246 :
1247 : /////////////////////////// Token.MISSING
1248 :
1249 :
1250 : static PyObject *_token_missing;
1251 :
1252 :
1253 : typedef struct {
1254 : PyObject_HEAD
1255 : } PyContextTokenMissing;
1256 :
1257 :
1258 : static PyObject *
1259 0 : context_token_missing_tp_repr(PyObject *self)
1260 : {
1261 0 : return PyUnicode_FromString("<Token.MISSING>");
1262 : }
1263 :
1264 :
1265 : PyTypeObject _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 :
1275 : static PyObject *
1276 2965 : get_token_missing(void)
1277 : {
1278 2965 : if (_token_missing != NULL) {
1279 2 : Py_INCREF(_token_missing);
1280 2 : return _token_missing;
1281 : }
1282 :
1283 2963 : _token_missing = (PyObject *)PyObject_New(
1284 : PyContextTokenMissing, &_PyContextTokenMissing_Type);
1285 2963 : if (_token_missing == NULL) {
1286 0 : return NULL;
1287 : }
1288 :
1289 2963 : Py_INCREF(_token_missing);
1290 2963 : return _token_missing;
1291 : }
1292 :
1293 :
1294 : ///////////////////////////
1295 :
1296 :
1297 : void
1298 30840 : _PyContext_ClearFreeList(PyInterpreterState *interp)
1299 : {
1300 : #if PyContext_MAXFREELIST > 0
1301 30840 : struct _Py_context_state *state = &interp->context;
1302 39106 : for (; state->numfree; state->numfree--) {
1303 8266 : PyContext *ctx = state->freelist;
1304 8266 : state->freelist = (PyContext *)ctx->ctx_weakreflist;
1305 8266 : ctx->ctx_weakreflist = NULL;
1306 8266 : PyObject_GC_Del(ctx);
1307 : }
1308 : #endif
1309 30840 : }
1310 :
1311 :
1312 : void
1313 3120 : _PyContext_Fini(PyInterpreterState *interp)
1314 : {
1315 3120 : if (_Py_IsMainInterpreter(interp)) {
1316 2951 : Py_CLEAR(_token_missing);
1317 : }
1318 3120 : _PyContext_ClearFreeList(interp);
1319 : #if defined(Py_DEBUG) && PyContext_MAXFREELIST > 0
1320 3120 : struct _Py_context_state *state = &interp->context;
1321 3120 : state->numfree = -1;
1322 : #endif
1323 3120 : _PyHamt_Fini(interp);
1324 3120 : }
1325 :
1326 :
1327 : PyStatus
1328 3134 : _PyContext_Init(PyInterpreterState *interp)
1329 : {
1330 3134 : if (!_Py_IsMainInterpreter(interp)) {
1331 171 : return _PyStatus_OK();
1332 : }
1333 :
1334 2963 : PyObject *missing = get_token_missing();
1335 2963 : if (PyDict_SetItemString(
1336 : PyContextToken_Type.tp_dict, "MISSING", missing))
1337 : {
1338 0 : Py_DECREF(missing);
1339 0 : return _PyStatus_ERR("can't init context types");
1340 : }
1341 2963 : Py_DECREF(missing);
1342 :
1343 2963 : return _PyStatus_OK();
1344 : }
|
