From: Ben Darnell <ben@bendarnell.com>
Date: Sun, 12 May 2013 17:17:20 +0000 (-0400)
Subject: Add a run_with_stack_context function to more easily use stack contexts in coroutines.
X-Git-Tag: v3.1.0~76^2~15
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=3305d1921ccf0d8c437cfe882e50ab1ad74a895a;p=thirdparty%2Ftornado.git

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 ec94a1b8e..642db1b35 100644
--- 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 50711ee1b..d85ea5071 100644
--- 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()