From 440e5355cfd0e3b1da5e5782568270d2ccd13dac Mon Sep 17 00:00:00 2001
From: Ben Darnell <ben@bendarnell.com>
Date: Sat, 8 Dec 2012 19:34:16 -0500
Subject: [PATCH] Replace a couple of contextmanager functions with
 ExceptionStackContext.

---
 tornado/simple_httpclient.py | 36 ++++++++++++++++--------------------
 tornado/testing.py           | 16 ++++++----------
 2 files changed, 22 insertions(+), 30 deletions(-)

diff --git a/tornado/simple_httpclient.py b/tornado/simple_httpclient.py
index 2c759a6d..7000d987 100644
--- a/tornado/simple_httpclient.py
+++ b/tornado/simple_httpclient.py
@@ -12,7 +12,6 @@ from tornado.util import b, GzipDecompressor
 
 import base64
 import collections
-import contextlib
 import copy
 import functools
 import os.path
@@ -134,7 +133,7 @@ class _HTTPConnection(object):
         self._decompressor = None
         # Timeout handle returned by IOLoop.add_timeout
         self._timeout = None
-        with stack_context.StackContext(self.cleanup):
+        with stack_context.ExceptionStackContext(self._handle_exception):
             self.parsed = urlparse.urlsplit(_unicode(self.request.url))
             if ssl is None and self.parsed.scheme == "https":
                 raise ValueError("HTTPS requires either python2.6+ or "
@@ -311,25 +310,22 @@ class _HTTPConnection(object):
             self.final_callback = None
             self.io_loop.add_callback(final_callback, response)
 
-    @contextlib.contextmanager
-    def cleanup(self):
-        try:
-            yield
-        except Exception, e:
-            if self.final_callback:
-                gen_log.warning("uncaught exception", exc_info=True)
-                self._run_callback(HTTPResponse(self.request, 599, error=e,
-                                    request_time=self.io_loop.time() - self.start_time,
-                                    ))
+    def _handle_exception(self, typ, value, tb):
+        if self.final_callback:
+            gen_log.warning("uncaught exception", exc_info=(typ, value, tb))
+            self._run_callback(HTTPResponse(self.request, 599, error=value,
+                                request_time=self.io_loop.time() - self.start_time,
+                                ))
 
-                if hasattr(self, "stream"):
-                    self.stream.close()
-            else:
-                # If our callback has already been called, we are probably
-                # catching an exception that is not caused by us but rather
-                # some child of our callback. Rather than drop it on the floor,
-                # pass it along.
-                raise
+            if hasattr(self, "stream"):
+                self.stream.close()
+            return True
+        else:
+            # If our callback has already been called, we are probably
+            # catching an exception that is not caused by us but rather
+            # some child of our callback. Rather than drop it on the floor,
+            # pass it along.
+            return False
 
     def _on_close(self):
         if self.final_callback is not None:
diff --git a/tornado/testing.py b/tornado/testing.py
index 59456433..22376627 100644
--- a/tornado/testing.py
+++ b/tornado/testing.py
@@ -36,9 +36,8 @@ except ImportError:
     netutil = None
     SimpleAsyncHTTPClient = None
 from tornado.log import gen_log
-from tornado.stack_context import StackContext
+from tornado.stack_context import ExceptionStackContext
 from tornado.util import raise_exc_info
-import contextlib
 import logging
 import os
 import re
@@ -167,13 +166,10 @@ class AsyncTestCase(unittest.TestCase):
         '''
         return IOLoop()
 
-    @contextlib.contextmanager
-    def _stack_context(self):
-        try:
-            yield
-        except Exception:
-            self.__failure = sys.exc_info()
-            self.stop()
+    def _handle_exception(self, typ, value, tb):
+        self.__failure = sys.exc_info()
+        self.stop()
+        return True
 
     def __rethrow(self):
         if self.__failure is not None:
@@ -182,7 +178,7 @@ class AsyncTestCase(unittest.TestCase):
             raise_exc_info(failure)
 
     def run(self, result=None):
-        with StackContext(self._stack_context):
+        with ExceptionStackContext(self._handle_exception):
             super(AsyncTestCase, self).run(result)
         # In case an exception escaped super.run or the StackContext caught
         # an exception when there wasn't a wait() to re-raise it, do so here.