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