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