1# -*- coding: utf-8 -*- 2""" 3Cooperative ``contextvars`` module. 4 5This module was added to Python 3.7. The gevent version is available 6on all supported versions of Python. However, see an important note 7about gevent 20.9. 8 9Context variables are like greenlet-local variables, just more 10inconvenient to use. They were designed to work around limitations in 11:mod:`asyncio` and are rarely needed by greenlet-based code. 12 13The primary difference is that snapshots of the state of all context 14variables in a given greenlet can be taken, and later restored for 15execution; modifications to context variables are "scoped" to the 16duration that a particular context is active. (This state-restoration 17support is rarely useful for greenlets because instead of always 18running "tasks" sequentially within a single thread like `asyncio` 19does, greenlet-based code usually spawns new greenlets to handle each 20task.) 21 22The gevent implementation is based on the Python reference implementation 23from :pep:`567` and doesn't have much optimization. In particular, setting 24context values isn't constant time. 25 26.. versionadded:: 1.5a3 27.. versionchanged:: 20.9.0 28 On Python 3.7 and above, this module is no longer monkey-patched 29 in place of the standard library version. 30 gevent depends on greenlet 0.4.17 which includes support for context variables. 31 This means that any number of greenlets can be running any number of asyncio tasks 32 each with their own context variables. This module is only greenlet aware, not 33 asyncio task aware, so its use is not recommended on Python 3.7 and above. 34 35 On previous versions of Python, this module continues to be a solution for 36 backporting code. It is also available if you wish to use the contextvar API 37 in a strictly greenlet-local manner. 38""" 39from __future__ import absolute_import 40from __future__ import division 41from __future__ import print_function 42 43 44__all__ = [ 45 'ContextVar', 46 'Context', 47 'copy_context', 48 'Token', 49] 50 51try: 52 from collections.abc import Mapping 53except ImportError: 54 from collections import Mapping # pylint:disable=deprecated-class 55 56from gevent._compat import PY37 57from gevent._util import _NONE 58from gevent.local import local 59 60__stdlib_expected__ = __all__ 61__implements__ = __stdlib_expected__ if PY37 else None 62 63# In the reference implementation, the interpreter level OS thread state 64# is modified to contain a pointer to the current context. Obviously we can't 65# touch that here because we're not tied to CPython's internals; plus, of course, 66# we want to operate with greenlets, not OS threads. So we use a greenlet-local object 67# to store the active context. 68class _ContextState(local): 69 70 def __init__(self): 71 self.context = Context() 72 73 74def _not_base_type(cls): 75 # This is not given in the PEP but is tested in test_context. 76 # Assign this method to __init_subclass__ in each type that can't 77 # be subclassed. (This only works in 3.6+, but context vars are only in 78 # 3.7+) 79 raise TypeError("not an acceptable base type") 80 81class _ContextData(object): 82 """ 83 A copy-on-write immutable mapping from ContextVar 84 keys to arbitrary values. Setting values requires a 85 copy, making it O(n), not O(1). 86 """ 87 88 # In theory, the HAMT used by the stdlib contextvars module could 89 # be used: It's often available at _testcapi.hamt() (see 90 # test_context). We'd need to be sure to add a correct __hash__ 91 # method to ContextVar to make that work well. (See 92 # Python/context.c:contextvar_generate_hash.) 93 94 __slots__ = ( 95 '_mapping', 96 ) 97 98 def __init__(self): 99 self._mapping = dict() 100 101 def __getitem__(self, key): 102 return self._mapping[key] 103 104 def __contains__(self, key): 105 return key in self._mapping 106 107 def __len__(self): 108 return len(self._mapping) 109 110 def __iter__(self): 111 return iter(self._mapping) 112 113 def set(self, key, value): 114 copy = _ContextData() 115 copy._mapping = self._mapping.copy() 116 copy._mapping[key] = value 117 return copy 118 119 def delete(self, key): 120 copy = _ContextData() 121 copy._mapping = self._mapping.copy() 122 del copy._mapping[key] 123 return copy 124 125 126class ContextVar(object): 127 """ 128 Implementation of :class:`contextvars.ContextVar`. 129 """ 130 131 __slots__ = ( 132 '_name', 133 '_default', 134 ) 135 136 def __init__(self, name, default=_NONE): 137 self._name = name 138 self._default = default 139 140 __init_subclass__ = classmethod(_not_base_type) 141 142 @classmethod 143 def __class_getitem__(cls, _): 144 # For typing support: ContextVar[str]. 145 # Not in the PEP. 146 # sigh. 147 return cls 148 149 @property 150 def name(self): 151 return self._name 152 153 def get(self, default=_NONE): 154 context = _context_state.context 155 try: 156 return context[self] 157 except KeyError: 158 pass 159 160 if default is not _NONE: 161 return default 162 163 if self._default is not _NONE: 164 return self._default 165 166 raise LookupError 167 168 def set(self, value): 169 context = _context_state.context 170 return context._set_value(self, value) 171 172 def reset(self, token): 173 token._reset(self) 174 175 def __repr__(self): 176 # This is not captured in the PEP but is tested by test_context 177 return '<%s.%s name=%r default=%r at 0x%x>' % ( 178 type(self).__module__, 179 type(self).__name__, 180 self._name, 181 self._default, 182 id(self) 183 ) 184 185 186class Token(object): 187 """ 188 Opaque implementation of :class:`contextvars.Token`. 189 """ 190 191 MISSING = _NONE 192 193 __slots__ = ( 194 '_context', 195 '_var', 196 '_old_value', 197 '_used', 198 ) 199 200 def __init__(self, context, var, old_value): 201 self._context = context 202 self._var = var 203 self._old_value = old_value 204 self._used = False 205 206 __init_subclass__ = classmethod(_not_base_type) 207 208 @property 209 def var(self): 210 """ 211 A read-only attribute pointing to the variable that created the token 212 """ 213 return self._var 214 215 @property 216 def old_value(self): 217 """ 218 A read-only attribute set to the value the variable had before 219 the ``set()`` call, or to :attr:`MISSING` if the variable wasn't set 220 before. 221 """ 222 return self._old_value 223 224 def _reset(self, var): 225 if self._used: 226 raise RuntimeError("Taken has already been used once") 227 228 if self._var is not var: 229 raise ValueError("Token was created by a different ContextVar") 230 231 if self._context is not _context_state.context: 232 raise ValueError("Token was created in a different Context") 233 234 self._used = True 235 if self._old_value is self.MISSING: 236 self._context._delete(var) 237 else: 238 self._context._reset_value(var, self._old_value) 239 240 def __repr__(self): 241 # This is not captured in the PEP but is tested by test_context 242 return '<%s.%s%s var=%r at 0x%x>' % ( 243 type(self).__module__, 244 type(self).__name__, 245 ' used' if self._used else '', 246 self._var, 247 id(self), 248 ) 249 250class Context(Mapping): 251 """ 252 Implementation of :class:`contextvars.Context` 253 """ 254 255 __slots__ = ( 256 '_data', 257 '_prev_context', 258 ) 259 260 def __init__(self): 261 """ 262 Creates an empty context. 263 """ 264 self._data = _ContextData() 265 self._prev_context = None 266 267 __init_subclass__ = classmethod(_not_base_type) 268 269 def run(self, function, *args, **kwargs): 270 if self._prev_context is not None: 271 raise RuntimeError( 272 "Cannot enter context; %s is already entered" % (self,) 273 ) 274 275 self._prev_context = _context_state.context 276 try: 277 _context_state.context = self 278 return function(*args, **kwargs) 279 finally: 280 _context_state.context = self._prev_context 281 self._prev_context = None 282 283 def copy(self): 284 """ 285 Return a shallow copy. 286 """ 287 result = Context() 288 result._data = self._data 289 return result 290 291 ### 292 # Operations used by ContextVar and Token 293 ### 294 295 def _set_value(self, var, value): 296 try: 297 old_value = self._data[var] 298 except KeyError: 299 old_value = Token.MISSING 300 301 self._data = self._data.set(var, value) 302 return Token(self, var, old_value) 303 304 def _delete(self, var): 305 self._data = self._data.delete(var) 306 307 def _reset_value(self, var, old_value): 308 self._data = self._data.set(var, old_value) 309 310 # Note that all Mapping methods, including Context.__getitem__ and 311 # Context.get, ignore default values for context variables (i.e. 312 # ContextVar.default). This means that for a variable var that was 313 # created with a default value and was not set in the context: 314 # 315 # - context[var] raises a KeyError, 316 # - var in context returns False, 317 # - the variable isn't included in context.items(), etc. 318 319 # Checking the type of key isn't part of the PEP but is tested by 320 # test_context.py. 321 @staticmethod 322 def __check_key(key): 323 if type(key) is not ContextVar: # pylint:disable=unidiomatic-typecheck 324 raise TypeError("ContextVar key was expected") 325 326 def __getitem__(self, key): 327 self.__check_key(key) 328 return self._data[key] 329 330 def __contains__(self, key): 331 self.__check_key(key) 332 return key in self._data 333 334 def __len__(self): 335 return len(self._data) 336 337 def __iter__(self): 338 return iter(self._data) 339 340 341def copy_context(): 342 """ 343 Return a shallow copy of the current context. 344 """ 345 return _context_state.context.copy() 346 347 348_context_state = _ContextState() 349