Task.after_return: Handler called after the task returns. Thanks to jonozzz. Closes #114.

celery/execute/trace.py

@@ -67,9 +67,18 @@ class TaskTrace(object):
         trace = TraceInfo.trace(self.task, self.args, self.kwargs)
         self.status = trace.status
         self.strtb = trace.strtb
+        self.handle_after_return(trace.status, trace.retval,
+                                 trace.exc_type, trace.tb, trace.strtb)
         handler = self._trace_handlers[trace.status]
         return handler(trace.retval, trace.exc_type, trace.tb, trace.strtb)
 
+    def handle_after_return(self, status, retval, type_, tb, strtb):
+        einfo = None
+        if status in states.EXCEPTION_STATES:
+            einfo = ExceptionInfo((retval, type_, tb))
+        self.task.after_return(status, retval, self.task_id,
+                               self.args, self.kwargs, einfo=einfo)
+
     def handle_success(self, retval, *args):
         """Handle successful execution."""
         self.task.on_success(retval, self.task_id, self.args, self.kwargs)
@@ -95,5 +104,5 @@ class TaskTrace(object):
         """Handle exception."""
         einfo = ExceptionInfo((type_, exc, tb))
         self.task.on_failure(exc, self.task_id,
-                             self.args, self.kwargs, einfo=info)
+                             self.args, self.kwargs, einfo=einfo)
         return einfo

celery/task/base.py

@@ -430,6 +430,23 @@ class Task(object):
         """
         pass
 
+    def after_return(self, status, retval, task_id, args, kwargs, einfo=None):
+        """Handler called after the task returns.
+
+        :param status: Current task state.
+        :param retval: Task return value/exception.
+        :param task_id: Unique id of the task.
+        :param args: Original arguments for the task that failed.
+        :param kwargs: Original keyword arguments for the task that failed.
+
+        :keyword einfo: :class:`celery.datastructures.ExceptionInfo` instance,
+           containing the traceback (if any).
+
+        The return value of this handler is ignored.
+
+        """
+        pass
+
     def on_failure(self, exc, task_id, args, kwargs, einfo=None):
         """Error handler.