Make StackContext more usable in libraries by reducing the need to

explicitly use IOLoop.add_callback or NullContext.

This change also simplifies and codifies the ordering behavior of
StackContext:  When a wrapped callback is executed, the innermost
StackContexts are those that were active when the callback was
wrapped.  Additional StackContexts (possibly including other copies
of the same contexts) may be present in outer levels of the stack, but
since the role of most StackContexts is to catch exceptions this should
not have any effect.

Backwards compatibility note:  StackContexts that do something other than
catch exceptions may need to be modified to be reentrant.

diff --git a/tornado/stack_context.py b/tornado/stack_context.py
index 8905ca62be730dee863777b3a777979aeb23edc8..9a835e7c7831df563dd0ec166d6881af130c82c2 100644 (file)
--- a/tornado/stack_context.py
+++ b/tornado/stack_context.py
@@ -104,23 +104,21 @@ def wrap(fn):
     # functools.wraps doesn't appear to work on functools.partial objects
     #@functools.wraps(fn)
     def wrapped(callback, contexts, *args, **kwargs):
-        # _state.contexts and contexts may share a common prefix.
-        # For each element of contexts not in that prefix, create a new
-        # StackContext object.
-        # TODO(bdarnell): do we want to be strict about the order,
-        # or is what we really want just set(contexts) - set(_state.contexts)?
-        # I think we do want to be strict about using identity comparison,
-        # so a set may not be quite right.  Conversely, it's not very stack-like
-        # to have new contexts pop up in the middle, so would we want to
-        # ensure there are no existing contexts not in the stack being restored?
-        # That feels right, but given the difficulty of handling errors at this
-        # level I'm not going to check for it now.
-        pairs = itertools.izip(itertools.chain(_state.contexts,
-                                               itertools.repeat(None)),
-                               contexts)
-        new_contexts = []
-        for old, new in itertools.dropwhile(lambda x: x[0] is x[1], pairs):
-            new_contexts.append(StackContext(new))
+        # If we're moving down the stack, _state.contexts is a prefix
+        # of contexts.  For each element of contexts not in that prefix,
+        # create a new StackContext object.
+        # If we're moving up the stack (or to an entirely different stack),
+        # _state.contexts will have elements not in contexts.  Use
+        # NullContext to clear the state and then recreate from contexts.
+        if (len(_state.contexts) > len(contexts) or
+            any(a is not b
+                for a, b in itertools.izip(_state.contexts, contexts))):
+            # contexts have been removed or changed, so start over
+            new_contexts = ([NullContext()] +
+                            [StackContext(c) for c in contexts])
+        else:
+            new_contexts = [StackContext(c)
+                            for c in contexts[len(_state.contexts):]]
         if new_contexts:
             with contextlib.nested(*new_contexts):
                 callback(*args, **kwargs)

diff --git a/tornado/test/stack_context_test.py b/tornado/test/stack_context_test.py
index 4237d344824911641a8dd1d5e6d7d5d6f5de7998..f6e7421bc0a65333623f90da0a611d0b74756e63 100755 (executable)
--- a/tornado/test/stack_context_test.py
+++ b/tornado/test/stack_context_test.py
@@ -55,15 +55,13 @@ class HTTPStackContextTest(AsyncHTTPTestCase, LogTrapTestCase):
 class StackContextTest(AsyncTestCase, LogTrapTestCase):
     def setUp(self):
         super(StackContextTest, self).setUp()
-        self.active_contexts = set()
+        self.active_contexts = []
 
     @contextlib.contextmanager
     def context(self, name):
-        assert name not in self.active_contexts
-        self.active_contexts.add(name)
+        self.active_contexts.append(name)
         yield
-        assert name in self.active_contexts
-        self.active_contexts.remove(name)
+        self.assertEqual(self.active_contexts.pop(), name)
 
     # Simulates the effect of an asynchronous library that uses its own
     # StackContext internally and then returns control to the application.
@@ -75,15 +73,15 @@ class StackContextTest(AsyncTestCase, LogTrapTestCase):
                 self.io_loop.add_callback(
                   functools.partial(library_inner_callback, callback))
         def library_inner_callback(callback):
-            assert 'application' in self.active_contexts
-            assert 'library' in self.active_contexts
-            # pass the callback out to the IOLoop to get out of the library
-            # context (could also use a NullContext here, but that would result
-            # in multiple instantiations of the application context)
-            self.io_loop.add_callback(callback)
+            self.assertEqual(self.active_contexts[-2:],
+                             ['application', 'library'])
+            callback()
         def final_callback():
-            assert 'application' in self.active_contexts
-            assert 'library' not in self.active_contexts
+            # implementation detail:  the full context stack at this point
+            # is ['application', 'library', 'application'].  The 'library'
+            # context was not removed, but is no longer innermost so
+            # the application context takes precedence.
+            self.assertEqual(self.active_contexts[-1], 'application')
             self.stop()
         with StackContext(functools.partial(self.context, 'application')):
             library_function(final_callback)