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

3 files changed, 44 insertions, 10 deletions

diff --git a/doc/changelog.rst b/doc/changelog.rst
index d402d60..2ef1c16 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -27,6 +27,8 @@ Major bugfixes:
 
 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
 * 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/futures.py b/trollius/futures.py
index 780c73d..3369c26 100644
--- a/trollius/futures.py
+++ b/trollius/futures.py
@@ -340,7 +340,21 @@ class Future(object):
         self._state = _FINISHED
         self._schedule_callbacks()
 
+    def _get_exception_tb(self):
+        """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]
+
     def set_exception(self, exception):
+        self._set_exception_with_tb(exception, None)
+
+    def _set_exception_with_tb(self, exception, exc_tb):
         """Mark the future done and set an exception.
 
         If the future is already done when this method is called, raises
@@ -367,7 +381,10 @@ class Future(object):
                 if self._loop.get_debug():
                     frame = sys._getframe(1)
                     tb = ['Traceback (most recent call last):\n']
-                    tb += traceback.format_stack(frame)
+                    if exc_tb:
+                        tb += exc_tb
+                    else:
+                        tb += traceback.format_tb(sys.exc_info()[2])
                     tb += traceback.format_exception_only(type(exception), exception)
                     self._tb_logger.tb = tb
                 else:
diff --git a/trollius/tasks.py b/trollius/tasks.py
index ab5c226..0915120 100644
--- a/trollius/tasks.py
+++ b/trollius/tasks.py
@@ -229,7 +229,7 @@ class Task(futures.Future):
         self._must_cancel = True
         return True
 
-    def _step(self, value=None, exc=None):
+    def _step(self, value=None, exc=None, exc_tb=None):
         assert not self.done(), \
             '_step(): already done: {0!r}, {1!r}, {2!r}'.format(self, value, exc)
         if self._must_cancel:
@@ -240,6 +240,7 @@ class Task(futures.Future):
         self._fut_waiter = None
 
         self.__class__._current_tasks[self._loop] = self
+        init_exc = exc
         # Call either coro.throw(exc) or coro.send(value).
         try:
             if exc is not None:
@@ -262,7 +263,10 @@ class Task(futures.Future):
         except futures.CancelledError as exc:
             super(Task, self).cancel()  # I.e., Future.cancel(self).
         except Exception as exc:
-            self.set_exception(exc)
+            if exc is init_exc:
+                self._set_exception_with_tb(exc, exc_tb)
+            else:
+                self.set_exception(exc)
         except BaseException as exc:
             self.set_exception(exc)
             raise
@@ -308,16 +312,27 @@ class Task(futures.Future):
                         'Task got bad yield: {0!r}'.format(result)))
         finally:
             self.__class__._current_tasks.pop(self._loop)
-            self = None  # Needed to break cycles when an exception occurs.
+
+            # Needed to break cycles when an exception occurs.
+            self = None
+            init_exc = None
 
     def _wakeup(self, future):
-        try:
-            value = future.result()
-        except Exception as exc:
-            # This may also be a cancellation.
-            self._step(None, exc)
+        if (future._state == futures._FINISHED
+        and future._exception is not None):
+            # Get the traceback before calling exception(), because the call
+            # to exception() clears the traceback
+            exc_tb = future._get_exception_tb()
+            exc = future.exception()
+            self._step(None, exc, exc_tb)
         else:
-            self._step(value, None)
+            try:
+                value = future.result()
+            except Exception as exc:
+                # This may also be a cancellation.
+                self._step(None, exc)
+            else:
+                self._step(value, None)
         self = None  # Needed to break cycles when an exception occurs.