4
4
import logging
5
5
import math
6
6
import os
7
+ import random
7
8
import re
8
9
import subprocess
9
10
import sys
@@ -1248,24 +1249,49 @@ def _make_threadlocal_contextvars(local):
1248
1249
class ContextVar (object ):
1249
1250
# Super-limited impl of ContextVar
1250
1251
1251
- def __init__ (self , name ):
1252
- # type: (str) -> None
1252
+ def __init__ (self , name , default = None ):
1253
+ # type: (str, Any ) -> None
1253
1254
self ._name = name
1255
+ self ._default = default
1254
1256
self ._local = local ()
1257
+ self ._original_local = local ()
1255
1258
1256
- def get (self , default ):
1259
+ def get (self , default = None ):
1257
1260
# type: (Any) -> Any
1258
- return getattr (self ._local , "value" , default )
1261
+ return getattr (self ._local , "value" , default or self . _default )
1259
1262
1260
1263
def set (self , value ):
1261
- # type: (Any) -> None
1264
+ # type: (Any) -> Any
1265
+ token = str (random .getrandbits (64 ))
1266
+ original_value = self .get ()
1267
+ setattr (self ._original_local , token , original_value )
1262
1268
self ._local .value = value
1269
+ return token
1270
+
1271
+ def reset (self , token ):
1272
+ # type: (Any) -> None
1273
+ self ._local .value = getattr (self ._original_local , token )
1274
+ del self ._original_local [token ]
1263
1275
1264
1276
return ContextVar
1265
1277
1266
1278
1279
+ def _make_noop_copy_context ():
1280
+ # type: () -> Callable[[], Any]
1281
+ class NoOpContext :
1282
+ def run (self , func , * args , ** kwargs ):
1283
+ # type: (Callable[..., Any], *Any, **Any) -> Any
1284
+ return func (* args , ** kwargs )
1285
+
1286
+ def copy_context ():
1287
+ # type: () -> NoOpContext
1288
+ return NoOpContext ()
1289
+
1290
+ return copy_context
1291
+
1292
+
1267
1293
def _get_contextvars ():
1268
- # type: () -> Tuple[bool, type]
1294
+ # type: () -> Tuple[bool, type, Callable[[], Any] ]
1269
1295
"""
1270
1296
Figure out the "right" contextvars installation to use. Returns a
1271
1297
`contextvars.ContextVar`-like class with a limited API.
@@ -1281,28 +1307,28 @@ def _get_contextvars():
1281
1307
# `aiocontextvars` is absolutely required for functional
1282
1308
# contextvars on Python 3.6.
1283
1309
try :
1284
- from aiocontextvars import ContextVar
1310
+ from aiocontextvars import ContextVar , copy_context
1285
1311
1286
- return True , ContextVar
1312
+ return True , ContextVar , copy_context
1287
1313
except ImportError :
1288
1314
pass
1289
1315
else :
1290
1316
# On Python 3.7 contextvars are functional.
1291
1317
try :
1292
- from contextvars import ContextVar
1318
+ from contextvars import ContextVar , copy_context
1293
1319
1294
- return True , ContextVar
1320
+ return True , ContextVar , copy_context
1295
1321
except ImportError :
1296
1322
pass
1297
1323
1298
1324
# Fall back to basic thread-local usage.
1299
1325
1300
1326
from threading import local
1301
1327
1302
- return False , _make_threadlocal_contextvars (local )
1328
+ return False , _make_threadlocal_contextvars (local ), _make_noop_copy_context ()
1303
1329
1304
1330
1305
- HAS_REAL_CONTEXTVARS , ContextVar = _get_contextvars ()
1331
+ HAS_REAL_CONTEXTVARS , ContextVar , copy_context = _get_contextvars ()
1306
1332
1307
1333
CONTEXTVARS_ERROR_MESSAGE = """
1308
1334
0 commit comments