[#9719] Support contextvars in coroutines

src/twisted/internet/defer.py

@@ -19,6 +19,7 @@
 from __future__ import division, absolute_import, print_function
 
 import attr
+import functools
 import traceback
 import types
 import warnings
@@ -28,11 +29,13 @@
 
 # Twisted imports
 from twisted.python.compat import cmp, comparable
-from twisted.python import lockfile, failure
+from twisted.python import lockfile, failure, reflect
 from twisted.logger import Logger
 from twisted.python.deprecate import warnAboutFunction, deprecated
 
 log = Logger()
+_contextvars = reflect.requireModule('contextvars')
+
 
 
 class AlreadyCalledError(Exception):
@@ -257,14 +260,14 @@ class Deferred:
 
     _chainedTo = None
 
-    def __init__(self, canceller=None):
+    def __init__(self, canceller=None, context=None):
         """
         Initialize a L{Deferred}.
 
         @param canceller: a callable used to stop the pending operation
-            scheduled by this L{Deferred} when L{Deferred.cancel} is
-            invoked. The canceller will be passed the deferred whose
-            cancelation is requested (i.e., self).
+            scheduled by this L{Deferred} when L{Deferred.cancel} is invoked.
+            The canceller will be passed the deferred whose cancelation is
+            requested (i.e., self).
 
             If a canceller is not given, or does not invoke its argument's
             C{callback} or C{errback} method, L{Deferred.cancel} will
@@ -280,13 +283,28 @@ def __init__(self, canceller=None):
 
         @type canceller: a 1-argument callable which takes a L{Deferred}. The
             return result is ignored.
+
+        @param context: The context that this Deferred should run its callbacks
+        under.
+
+        @type context: L{contextvars.Context}
         """
         self.callbacks = []
         self._canceller = canceller
         if self.debug:
             self._debugInfo = DebugInfo()
             self._debugInfo.creator = traceback.format_stack()[:-1]
 
+        if _contextvars and context is None:
+            self._context = _contextvars.copy_context()
+        elif _contextvars is None and context:
+            raise Exception(
+                "Manually supplying a contextvars context is not supported "
+                "if contextvars is not installed."
+            )
+        else:
+            self._context = context
+
 
     def addCallbacks(self, callback, errback=None,
                      callbackArgs=None, callbackKeywords=None,
@@ -649,10 +667,17 @@ def _runCallbacks(self):
                 try:
                     current._runningCallbacks = True
                     try:
+                        originalCallback = callback
+                        if _contextvars:
+                            callback = functools.partial(
+                                current._context.run, callback
+                            )
+
                         current.result = callback(current.result, *args, **kw)
+
                         if current.result is current:
                             warnAboutFunction(
-                                callback,
+                                originalCallback,
                                 "Callback returned the Deferred "
                                 "it was attached to; this breaks the "
                                 "callback chain and will raise an "
@@ -1400,17 +1425,29 @@ def _inlineCallbacks(result, g, status):
     # loop and the waiting variable solve that by manually unfolding the
     # recursion.
 
-    waiting = [True, # waiting for result?
-               None] # result
+    waiting = [True,  # waiting for result?
+               None]  # result
+
+    if _contextvars:
+        current_context = _contextvars.copy_context()
 
     while 1:
         try:
             # Send the last result back as the result of the yield expression.
             isFailure = isinstance(result, failure.Failure)
-            if isFailure:
-                result = result.throwExceptionIntoGenerator(g)
+
+            if _contextvars:
+                if isFailure:
+                    result = current_context.run(
+                        result.throwExceptionIntoGenerator, g
+                    )
+                else:
+                    result = current_context.run(g.send, result)
             else:
-                result = g.send(result)
+                if isFailure:
+                    result = result.throwExceptionIntoGenerator(g)
+                else:
+                    result = g.send(result)
         except StopIteration as e:
             # fell off the end, or "return" statement
             status.deferred.callback(getattr(e, "value", None))
@@ -1466,8 +1503,11 @@ def gotResult(r):
                     waiting[0] = False
                     waiting[1] = r
                 else:
-                    # We are not waiting for deferred result any more
-                    _inlineCallbacks(r, g, status)
+                    if _contextvars:
+                        current_context.run(_inlineCallbacks, r, g, status)
+                    else:
+                        # We are not waiting for deferred result any more
+                        _inlineCallbacks(r, g, status)
 
             result.addBoth(gotResult)
             if waiting[0]:
@@ -2002,13 +2042,11 @@ def _tryLock():
 
 
 
-__all__ = ["Deferred", "DeferredList", "succeed", "fail", "FAILURE", "SUCCESS",
-           "AlreadyCalledError", "TimeoutError", "gatherResults",
-           "maybeDeferred", "ensureDeferred",
-           "waitForDeferred", "deferredGenerator", "inlineCallbacks",
-           "returnValue",
-           "DeferredLock", "DeferredSemaphore", "DeferredQueue",
-           "DeferredFilesystemLock", "AlreadyTryingToLockError",
-           "CancelledError",
-          ]
-
+__all__ = [
+    "Deferred", "DeferredList", "succeed", "fail", "FAILURE", "SUCCESS",
+    "AlreadyCalledError", "TimeoutError", "gatherResults",
+    "maybeDeferred", "ensureDeferred",
+    "waitForDeferred", "deferredGenerator", "inlineCallbacks", "returnValue",
+    "DeferredLock", "DeferredSemaphore", "DeferredQueue",
+    "DeferredFilesystemLock", "AlreadyTryingToLockError", "CancelledError",
+]

src/twisted/newsfragments/9719.feature

@@ -0,0 +1 @@
+twisted.internet.defer.Deferred will now capture the current contextvars context when created, and run callbacks using that context. This functionality requires Python 3.7+, or the contextvars PyPI backport to be installed for Python 3.5-3.6.