Add a run_with_stack_context function to more easily use stack contexts in coroutines.

diff --git a/tornado/stack_context.py b/tornado/stack_context.py
index ec94a1b8e1ae6c735f93aeeca3bbf715d10bbdff..642db1b35fd6e38430880f625f5cc5bd978cda1e 100644 (file)
--- a/tornado/stack_context.py
+++ b/tornado/stack_context.py
@@ -303,3 +303,27 @@ def _handle_exception(tail, exc):
         tail = tail.old_contexts[1]
 
     return exc
+
+
+def run_with_stack_context(context, func):
+    """Run a coroutine ``func`` in the given `StackContext`.
+
+    It is not safe to have a ``yield`` statement within a ``with StackContext``
+    block, so it is difficult to use stack context with `.gen.coroutine`.
+    This helper function runs the function in the correct context while
+    keeping the ``yield`` and ``with`` statements syntactically separate.
+
+    Example::
+
+        @gen.coroutine
+        def incorrect():
+            with StackContext(ctx):
+                # ERROR: this will raise StackContextInconsistentError
+                yield other_coroutine()
+
+        @gen.coroutine
+        def correct():
+            yield run_with_stack_context(StackContext(ctx), other_coroutine)
+    """
+    with context:
+        return func()

diff --git a/tornado/test/stack_context_test.py b/tornado/test/stack_context_test.py
index 50711ee1b39bd6532762ec8b7a9db2ff30d614a1..d85ea50719866d700c37f29ceaa9eb050b3db4f5 100644 (file)
--- a/tornado/test/stack_context_test.py
+++ b/tornado/test/stack_context_test.py
@@ -3,7 +3,7 @@ from __future__ import absolute_import, division, print_function, with_statement
 
 from tornado import gen
 from tornado.log import app_log
-from tornado.stack_context import StackContext, wrap, NullContext, StackContextInconsistentError, ExceptionStackContext
+from tornado.stack_context import StackContext, wrap, NullContext, StackContextInconsistentError, ExceptionStackContext, run_with_stack_context
 from tornado.testing import AsyncHTTPTestCase, AsyncTestCase, ExpectLog, gen_test
 from tornado.test.util import unittest
 from tornado.web import asynchronous, Application, RequestHandler
@@ -177,6 +177,26 @@ class StackContextTest(AsyncTestCase):
             self.io_loop.add_callback(cb)
         yield gen.Wait('k1')
 
+    def test_run_with_stack_context(self):
+        @gen.coroutine
+        def f1():
+            self.assertEqual(self.active_contexts, ['c1'])
+            yield run_with_stack_context(
+                StackContext(functools.partial(self.context, 'c1')),
+                f2)
+            self.assertEqual(self.active_contexts, ['c1'])
+
+        @gen.coroutine
+        def f2():
+            self.assertEqual(self.active_contexts, ['c1', 'c2'])
+            yield gen.Task(self.io_loop.add_callback)
+            self.assertEqual(self.active_contexts, ['c1', 'c2'])
+
+        self.assertEqual(self.active_contexts, [])
+        run_with_stack_context(
+            StackContext(functools.partial(self.context, 'c1')),
+            f1)
+        self.assertEqual(self.active_contexts, [])
 
 if __name__ == '__main__':
     unittest.main()