Add an implementation of contextvars.

These are greenlet local.

Fixes #1407

CHANGES.rst

@@ -54,6 +54,10 @@ Other
   closing files. Previously if a timeout error happened, a second call
   to ``communicate`` might not close the pipe.
 
+- Add `gevent.contextvars`, a cooperative version of `contextvars`.
+  This is available to all Python versions. On Python 3.7 and above,
+  where `contextvars` is a standard library module, it is
+  monkey-patched by default. See :issue:`1407`.
 
 1.5a3 (2020-01-01)
 ==================

docs/api/gevent.contextvars.rst

+==========================================================
+ :mod:`gevent.contextvars` -- Cooperative ``contextvars``
+==========================================================
+
+.. automodule:: gevent.contextvars
+    :members:

docs/api/index.rst

@@ -28,6 +28,7 @@ Synchronization primitives (locks, queues, events)
 
    gevent.event
    gevent.queue
+   gevent.local
    gevent.lock
 
 Low-level details
@@ -48,6 +49,7 @@ Module Listing
    gevent.backdoor
    gevent.baseserver
    gevent.builtins
+   gevent.contextvars
    gevent.core
    gevent.event
    gevent.events

src/gevent/_patcher.py

@@ -34,6 +34,7 @@ MAPPING = {
     'gevent.signal': 'signal',
     'gevent.time': 'time',
     'gevent.queue': 'queue' if PY3 else 'Queue',
+    'gevent.contextvars': 'contextvars',
 }
 
 _PATCH_PREFIX = '__g_patched_module_'

src/gevent/contextvars.py

+# -*- 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
+
+__implements__ = __all__ 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()

src/gevent/monkey.py

@@ -494,6 +494,19 @@ def patch_time():
     """
     _patch_module('time')
 
+@_ignores_DoNotPatch
+def patch_contextvars():
+    """
+    On Python 3.7 and above, replaces the implementations of :mod:`contextvars`
+    with :mod:`gevent.contextvars`.
+    """
+    try:
+        __import__('contextvars')
+    except ImportError:
+        pass
+    else:
+        _patch_module('contextvars')
+
 
 def _patch_existing_locks(threading):
     if len(list(threading.enumerate())) != 1:
@@ -1065,7 +1078,7 @@ def _subscribe_signal_os(will_patch_all):
 def patch_all(socket=True, dns=True, time=True, select=True, thread=True, os=True, ssl=True,
               subprocess=True, sys=False, aggressive=True, Event=True,
               builtins=True, signal=True,
-              queue=True,
+              queue=True, contextvars=True,
               **kwargs):
     """
     Do all of the default monkey patching (calls every other applicable
@@ -1094,6 +1107,8 @@ def patch_all(socket=True, dns=True, time=True, select=True, thread=True, os=Tru
        Add *queue*, defaulting to True, for Python 3.7.
     .. versionchanged:: 1.5
        Remove the ``httplib`` argument. Previously, setting it raised a ``ValueError``.
+    .. versionchanged:: 1.5a3
+       Add the ``contextvars`` argument.
     .. versionchanged:: 1.5
        Better handling of patching more than once.
     """
@@ -1143,6 +1158,8 @@ def patch_all(socket=True, dns=True, time=True, select=True, thread=True, os=Tru
         patch_signal()
     if queue:
         patch_queue()
+    if contextvars:
+        patch_contextvars()
 
     _notify_patch(events.GeventDidPatchBuiltinModulesEvent(modules_to_patch, kwargs), _warnings)
     _notify_patch(events.GeventDidPatchAllEvent(modules_to_patch, kwargs), _warnings)

src/gevent/testing/patched_tests_setup.py

@@ -230,6 +230,12 @@ disabled_tests = [
     # and this doesn't work reliably on all versions.
     'test_httplib.HeaderTests.test_headers_debuglevel',
 
+    # These depend on the exact error message produced by the interpreter
+    # when too many arguments are passed to functions. We can't match
+    # the exact signatures (because Python 2 doesn't support the syntax)
+    'test_context.ContextTest.test_context_new_1',
+    'test_context.ContextTest.test_context_var_new_1',
+
 ]
 
 if 'thread' in os.getenv('GEVENT_FILE', ''):