You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

337 lines
8.9 KiB
Python

5 years ago
# -*- coding: utf-8 -*-
"""
Cooperative ``contextvars`` module.
This module was added to Python 3.7. The gevent version is available
on all supported versions of Python.
Context variables are like greenlet-local variables, just more
inconvenient to use. They were designed to work around limitations in
:mod:`asyncio` and are rarely needed by greenlet-based code.
The primary difference is that snapshots of the state of all context
variables in a given greenlet can be taken, and later restored for
execution; modifications to context variables are "scoped" to the
duration that a particular context is active. (This state-restoration
support is rarely useful for greenlets because instead of always
running "tasks" sequentially within a single thread like `asyncio`
does, greenlet-based code usually spawns new greenlets to handle each
task.)
The gevent implementation is based on the Python reference implementation
from :pep:`567` and doesn't have much optimization. In particular, setting
context values isn't constant time.
.. versionadded:: 1.5a3
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
__all__ = [
'ContextVar',
'Context',
'copy_context',
'Token',
]
try:
from collections.abc import Mapping
except ImportError:
from collections import Mapping
from gevent._compat import PY37
from gevent._util import _NONE
from gevent.local import local
__stdlib_expected__ = __all__
__implements__ = __stdlib_expected__ if PY37 else None
# In the reference implementation, the interpreter level OS thread state
# is modified to contain a pointer to the current context. Obviously we can't
# touch that here because we're not tied to CPython's internals; plus, of course,
# we want to operate with greenlets, not OS threads. So we use a greenlet-local object
# to store the active context.
class _ContextState(local):
def __init__(self):
self.context = Context()
def _not_base_type(cls):
# This is not given in the PEP but is tested in test_context.
# Assign this method to __init_subclass__ in each type that can't
# be subclassed. (This only works in 3.6+, but context vars are only in
# 3.7+)
raise TypeError("not an acceptable base type")
class _ContextData(object):
"""
A copy-on-write immutable mapping from ContextVar
keys to arbitrary values. Setting values requires a
copy, making it O(n), not O(1).
"""
# In theory, the HAMT used by the stdlib contextvars module could
# be used: It's often available at _testcapi.hamt() (see
# test_context). We'd need to be sure to add a correct __hash__
# method to ContextVar to make that work well. (See
# Python/context.c:contextvar_generate_hash.)
__slots__ = (
'_mapping',
)
def __init__(self):
self._mapping = dict()
def __getitem__(self, key):
return self._mapping[key]
def __contains__(self, key):
return key in self._mapping
def __len__(self):
return len(self._mapping)
def __iter__(self):
return iter(self._mapping)
def set(self, key, value):
copy = _ContextData()
copy._mapping = self._mapping.copy()
copy._mapping[key] = value
return copy
def delete(self, key):
copy = _ContextData()
copy._mapping = self._mapping.copy()
del copy._mapping[key]
return copy
class ContextVar(object):
"""
Implementation of :class:`contextvars.ContextVar`.
"""
__slots__ = (
'_name',
'_default',
)
def __init__(self, name, default=_NONE):
self._name = name
self._default = default
__init_subclass__ = classmethod(_not_base_type)
@classmethod
def __class_getitem__(cls, _):
# For typing support: ContextVar[str].
# Not in the PEP.
# sigh.
return cls
@property
def name(self):
return self._name
def get(self, default=_NONE):
context = _context_state.context
try:
return context[self]
except KeyError:
pass
if default is not _NONE:
return default
if self._default is not _NONE:
return self._default
raise LookupError
def set(self, value):
context = _context_state.context
return context._set_value(self, value)
def reset(self, token):
token._reset(self)
def __repr__(self):
# This is not captured in the PEP but is tested by test_context
return '<%s.%s name=%r default=%r at 0x%x>' % (
type(self).__module__,
type(self).__name__,
self._name,
self._default,
id(self)
)
class Token(object):
"""
Opaque implementation of :class:`contextvars.Token`.
"""
MISSING = _NONE
__slots__ = (
'_context',
'_var',
'_old_value',
'_used',
)
def __init__(self, context, var, old_value):
self._context = context
self._var = var
self._old_value = old_value
self._used = False
__init_subclass__ = classmethod(_not_base_type)
@property
def var(self):
"""
A read-only attribute pointing to the variable that created the token
"""
return self._var
@property
def old_value(self):
"""
A read-only attribute set to the value the variable had before
the ``set()`` call, or to :attr:`MISSING` if the variable wasn't set
before.
"""
return self._old_value
def _reset(self, var):
if self._used:
raise RuntimeError("Taken has already been used once")
if self._var is not var:
raise ValueError("Token was created by a different ContextVar")
if self._context is not _context_state.context:
raise ValueError("Token was created in a different Context")
self._used = True
if self._old_value is self.MISSING:
self._context._delete(var)
else:
self._context._reset_value(var, self._old_value)
def __repr__(self):
# This is not captured in the PEP but is tested by test_context
return '<%s.%s%s var=%r at 0x%x>' % (
type(self).__module__,
type(self).__name__,
' used' if self._used else '',
self._var,
id(self),
)
class Context(Mapping):
"""
Implementation of :class:`contextvars.Context`
"""
__slots__ = (
'_data',
'_prev_context',
)
def __init__(self):
"""
Creates an empty context.
"""
self._data = _ContextData()
self._prev_context = None
__init_subclass__ = classmethod(_not_base_type)
def run(self, function, *args, **kwargs):
if self._prev_context is not None:
raise RuntimeError(
"Cannot enter context; %s is already entered" % (self,)
)
self._prev_context = _context_state.context
try:
_context_state.context = self
return function(*args, **kwargs)
finally:
_context_state.context = self._prev_context
self._prev_context = None
def copy(self):
"""
Return a shallow copy.
"""
result = Context()
result._data = self._data
return result
###
# Operations used by ContextVar and Token
###
def _set_value(self, var, value):
try:
old_value = self._data[var]
except KeyError:
old_value = Token.MISSING
self._data = self._data.set(var, value)
return Token(self, var, old_value)
def _delete(self, var):
self._data = self._data.delete(var)
def _reset_value(self, var, old_value):
self._data = self._data.set(var, old_value)
# Note that all Mapping methods, including Context.__getitem__ and
# Context.get, ignore default values for context variables (i.e.
# ContextVar.default). This means that for a variable var that was
# created with a default value and was not set in the context:
#
# - context[var] raises a KeyError,
# - var in context returns False,
# - the variable isn't included in context.items(), etc.
# Checking the type of key isn't part of the PEP but is tested by
# test_context.py.
@staticmethod
def __check_key(key):
if type(key) is not ContextVar: # pylint:disable=unidiomatic-typecheck
raise TypeError("ContextVar key was expected")
def __getitem__(self, key):
self.__check_key(key)
return self._data[key]
def __contains__(self, key):
self.__check_key(key)
return key in self._data
def __len__(self):
return len(self._data)
def __iter__(self):
return iter(self._data)
def copy_context():
"""
Return a shallow copy of the current context.
"""
return _context_state.context.copy()
_context_state = _ContextState()