in debug mode on python 2, store the exception traceback

4 files changed, 65 insertions, 12 deletions

diff --git a/doc/changelog.rst b/doc/changelog.rst
index 6a015dc..d64e47b 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -7,6 +7,11 @@ Version 1.0.3
 
 Changes:
 
+* On Python 2 in debug mode, Future.set_exception() now stores the traceback
+  object of the exception in addition to the exception object. When a task
+  waiting for another task and the other task raises an exception, the
+  traceback object is now copied with the exception. Be careful, storing the
+  traceback object may create reference leaks.
 * On Python 3.5 and newer, reuse socket.socketpair() in the windows_utils
   submodule.
 * On Python 3.4 and newer, use os.set_inheritable().
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 780c73d..02514f2 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
+        exc_tb = self._exception_tb
+        self._exception_tb = None
         if self._exception is not None:
-            raise self._exception
+            if exc_tb is not None:
+                compat.reraise(type(self._exception), self._exception, exc_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):
@@ -340,7 +351,13 @@ class Future(object):
         self._state = _FINISHED
         self._schedule_callbacks()
 
+    def _get_exception_tb(self):
+        return self._exception_tb
+
     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
@@ -351,6 +368,11 @@ class Future(object):
         if isinstance(exception, type):
             exception = exception()
         self._exception = exception
+        if exc_tb is not None:
+            self._exception_tb = exc_tb
+            exc_tb = None
+        elif self._loop.get_debug() and not compat.PY3:
+            self._exception_tb = sys.exc_info()[2]
         self._state = _FINISHED
         self._schedule_callbacks()
         if _PY34:
diff --git a/trollius/tasks.py b/trollius/tasks.py
index 892bb7c..428d3af 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:
@@ -239,6 +239,10 @@ class Task(futures.Future):
         coro = self._coro
         self._fut_waiter = None
 
+        if exc_tb is not None:
+            init_exc = exc
+        else:
+            init_exc = None
         self.__class__._current_tasks[self._loop] = self
         # Call either coro.throw(exc) or coro.send(value).
         try:
@@ -261,11 +265,16 @@ class Task(futures.Future):
             self.set_result(result)
         except futures.CancelledError as exc:
             super(Task, self).cancel()  # I.e., Future.cancel(self).
-        except Exception as exc:
-            self.set_exception(exc)
         except BaseException as exc:
-            self.set_exception(exc)
-            raise
+            if exc is init_exc:
+                self._set_exception_with_tb(exc, exc_tb)
+                exc_tb = None
+            else:
+                self.set_exception(exc)
+
+            if not isinstance(exc, Exception):
+                # reraise BaseException
+                raise
         else:
             if coroutines._DEBUG:
                 if not coroutines._coroutine_at_yield_from(self._coro):
@@ -313,13 +322,22 @@ class Task(futures.Future):
             self = None  # Needed to break cycles when an exception occurs.
 
     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 calling
+            # the exception() method clears the traceback
+            exc_tb = future._get_exception_tb()
+            exc = future.exception()
+            self._step(None, exc, exc_tb)
+            exc_tb = None
         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.