On Python 2, Future.exception() now raises the traceback with the original

traceback.

3 files changed, 29 insertions, 12 deletions

diff --git a/doc/changelog.rst b/doc/changelog.rst
index 2ef1c16..407ea60 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -29,6 +29,8 @@ Other changes:
 
 * On Python 2, Task waiting on a future now keeps the exception traceback of
   the future, instead of showing where the exception was re-raised
+* On Python 2, Future.exception() now raises the traceback with the original
+  traceback.
 * Return destructor logs the source traceback in debug mode
 * Python issue #22448: Improve cancelled timer callback handles cleanup. Patch
   by Joshua Moore-Oliva.
diff --git a/trollius/compat.py b/trollius/compat.py
index b93070d..7947842 100644
--- a/trollius/compat.py
+++ b/trollius/compat.py
@@ -51,3 +51,11 @@ def flatten_bytes(data):
         return data.tobytes()
     else:
         return data
+
+if PY3:
+    def reraise(tp, value, tb=None):
+        if value.__traceback__ is not tb:
+            raise value.with_traceback(tb)
+        raise value
+else:
+    exec("""def reraise(tp, value, tb=None):  raise tp, value, tb""")
diff --git a/trollius/futures.py b/trollius/futures.py
index 3369c26..68beb99 100644
--- a/trollius/futures.py
+++ b/trollius/futures.py
@@ -13,6 +13,7 @@ try:
 except ImportError:
     import repr as reprlib   # Python 2
 
+from . import compat
 from . import events
 from . import executor
 
@@ -137,6 +138,10 @@ class Future(object):
     _exception = None
     _loop = None
 
+    # Used by Python 2 to raise the exception with the original traceback
+    # in the exception() method
+    _exception_tb = None
+
     _log_traceback = False   # Used for Python >= 3.4
     _tb_logger = None        # Used for Python <= 3.3
 
@@ -273,8 +278,13 @@ class Future(object):
         if self._tb_logger is not None:
             self._tb_logger.clear()
             self._tb_logger = None
+        tb = self._exception_tb
+        self._exception_tb = None
         if self._exception is not None:
-            raise self._exception
+            if tb is not None:
+                compat.reraise(type(self._exception), self._exception, tb)
+            else:
+                raise self._exception
         return self._result
 
     def exception(self):
@@ -293,6 +303,7 @@ class Future(object):
         if self._tb_logger is not None:
             self._tb_logger.clear()
             self._tb_logger = None
+        self._exception_tb = None
         return self._exception
 
     def add_done_callback(self, fn):
@@ -344,12 +355,7 @@ class Future(object):
         """Helper method to call _set_exception_with_tb().
 
         Use it to get the traceback of a future to copy it to a new future."""
-        if _PY34:
-            return None
-        if self._tb_logger is None or not self._tb_logger.tb:
-            return None
-        # Ignore first and last line
-        return self._tb_logger.tb[1:-1]
+        return self._exception_tb
 
     def set_exception(self, exception):
         self._set_exception_with_tb(exception, None)
@@ -365,13 +371,17 @@ class Future(object):
         if isinstance(exception, type):
             exception = exception()
         self._exception = exception
+        if exc_tb is not None:
+            self._exception_tb = exc_tb
+        elif not compat.PY3:
+            self._exception_tb = sys.exc_info()[2]
         self._state = _FINISHED
         self._schedule_callbacks()
         if _PY34:
             self._log_traceback = True
         else:
             self._tb_logger = _TracebackLogger(self, exception)
-            if hasattr(exception, '__traceback__'):
+            if compat.PY3:
                 # Python 3: exception contains a link to the traceback
 
                 # Arrange for the logger to be activated after all callbacks
@@ -381,10 +391,7 @@ class Future(object):
                 if self._loop.get_debug():
                     frame = sys._getframe(1)
                     tb = ['Traceback (most recent call last):\n']
-                    if exc_tb:
-                        tb += exc_tb
-                    else:
-                        tb += traceback.format_tb(sys.exc_info()[2])
+                    tb += traceback.format_tb(self._exception_tb)
                     tb += traceback.format_exception_only(type(exception), exception)
                     self._tb_logger.tb = tb
                 else: